diff --git a/.gitignore b/.gitignore index a045926..2f81438 100644 --- a/.gitignore +++ b/.gitignore @@ -87,6 +87,7 @@ create.sql shell.nix **/.jekyll-cache **/tmp +**/*.bak # Helm kubernetes/packages diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index e455922..18730a7 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -309,8 +309,10 @@ python -m pytest tests/unit/test_utils.py python -m pytest tests/unit/test_utils.py::TestToJson::test_serializes_dict ``` -The authentication tests use Moto to emulate Cognito. Unit tests must not -require real AWS credentials or network access. +The authentication tests use Moto to emulate Cognito. ECS-spawner unit tests +use botocore's `Stubber` to validate ECS API requests and responses without AWS +credentials or network access; they do not start containers. A future +Docker-enabled LocalStack suite should cover ECS task image execution. ### UI unit tests diff --git a/helm/beakerhub/templates/beakerhub-config.yaml b/helm/beakerhub/templates/beakerhub-config.yaml index 79d5815..b0d8746 100644 --- a/helm/beakerhub/templates/beakerhub-config.yaml +++ b/helm/beakerhub/templates/beakerhub-config.yaml @@ -191,19 +191,17 @@ c.Authenticator.admin_users = {{ .Values.auth.adminUsers | toJson }} # Task Configuration (node image import jobs, etc.) #------------------------------------------------------------------------------ -c.BeakerHub.task_namespace = {{ .Values.namespace | quote }} -c.BeakerHub.task_reporter_image = {{ printf "%s%s:%s" (default "" .Values.defaultRegistry) .Values.tasks.reporter.image.repository .Values.tasks.reporter.image.tag | quote }} -c.BeakerHub.task_reporter_pull_policy = {{ .Values.tasks.reporter.image.pullPolicy | quote }} -c.BeakerHub.task_reporter_resources = {{ .Values.tasks.reporter.resources | toJson }} -c.BeakerHub.task_node_image_resources = {{ .Values.tasks.nodeImageDefaults.resources | toJson }} -c.BeakerHub.task_backoff_limit = {{ .Values.tasks.backoffLimit }} -c.BeakerHub.task_active_deadline_seconds = {{ .Values.tasks.activeDeadlineSeconds }} -c.BeakerHub.task_ttl_seconds_after_finished = {{ .Values.tasks.ttlSecondsAfterFinished }} +c.KubernetesTaskRunnerService.namespace = {{ .Values.namespace | quote }} +c.KubernetesTaskRunnerService.node_image_resources = {{ .Values.tasks.nodeImageDefaults.resources | toJson }} +c.KubernetesTaskRunnerService.backoff_limit = {{ .Values.tasks.backoffLimit }} +c.KubernetesTaskRunnerService.active_deadline_seconds = {{ .Values.tasks.activeDeadlineSeconds }} +c.KubernetesTaskRunnerService.ttl_seconds_after_finished = {{ .Values.tasks.ttlSecondsAfterFinished }} +c.KubernetesDashboardService.namespace = {{ .Values.namespace | quote }} {{- if .Values.tasks.nodeSelector }} -c.BeakerHub.task_node_selector = {{ .Values.tasks.nodeSelector | toJson }} +c.KubernetesTaskRunnerService.node_selector = {{ .Values.tasks.nodeSelector | toJson }} {{- end }} {{- if .Values.tasks.tolerations }} -c.BeakerHub.task_tolerations = {{ .Values.tasks.tolerations | toJson }} +c.KubernetesTaskRunnerService.tolerations = {{ .Values.tasks.tolerations | toJson }} {{- end }} {{- if .Values.config.extraConfig }} diff --git a/helm/beakerhub/templates/configmap.yaml b/helm/beakerhub/templates/configmap.yaml index a95b329..dd2b96f 100644 --- a/helm/beakerhub/templates/configmap.yaml +++ b/helm/beakerhub/templates/configmap.yaml @@ -36,7 +36,7 @@ data: "Content-Security-Policy": "frame-ancestors 'self' {{ .Values.ingress.hosts.proxy }} *.{{ .Values.ingress.hosts.proxy }}", } } - c.KernelProvisionerFactory.default_provisioner_name = "kubernetes-local-provisioner" + c.KernelProvisionerFactory.default_provisioner_name = "beakerhub-local-provisioner" c.BaseBeakerApp.secrets_manager_class = "beakerhub.services.secrets.beakerhub.BeakerhubSecretsManager" c.BeakerhubSecretsManager.policy_override_env_key_suffix = "{{ .Values.secrets.policyOverrideEnvSuffix }}" diff --git a/implicit-kubernetes-backend-assumptions.md b/implicit-kubernetes-backend-assumptions.md new file mode 100644 index 0000000..c63122b --- /dev/null +++ b/implicit-kubernetes-backend-assumptions.md @@ -0,0 +1,644 @@ +# Implicit Kubernetes Backend Assumptions + +This document maps backend code that assumes Kubernetes even when it is not +part of the configured spawner implementation. These assumptions need separate +handling if BeakerHub uses ECS as a runtime backend. + +## Scope + +Included: + +- Python backend code under `src/beakerhub` that directly uses Kubernetes + concepts or APIs. +- Backend data models and API fields whose names expose Kubernetes concepts. +- Node configuration that selects a Kubernetes-named implementation outside + the spawner itself. + +Excluded: + +- `src/beakerhub/spawner/kubernetes.py`, because it is the explicit Kubernetes + spawner implementation. +- `src/beakerhub/spawner/aws/ecs.py`, because it is the explicit ECS spawner + implementation. +- `helm/`, because Helm and its chart templates are Kubernetes deployment + artifacts by definition. + +No source file currently refers to EKS explicitly. The assumptions below are +generic Kubernetes assumptions. + +## Summary + +There are four significant areas of implicit coupling: + +1. `BeakerHub` selects the Kubernetes spawner by default and defines + Kubernetes-specific node-image task configuration. +2. Node-image imports are implemented directly as Kubernetes Jobs. +3. The admin dashboard and session-log endpoint query Kubernetes resources + directly. +4. Task persistence and APIs expose the Kubernetes term `job_name` as the + external runtime identifier. + +The node-side local provisioner also has a Kubernetes-specific name, although +its behavior is not inherently Kubernetes-specific. + +## Application defaults and configuration + +### Default spawner + +`src/beakerhub/app.py:126-129` imports and returns `BeakerKubeSpawner` from the +application's `spawner_class` default. + +This makes Kubernetes the implicit runtime whenever deployment configuration +does not explicitly select another spawner. The default is separate from the +Kubernetes spawner implementation itself and is therefore part of backend +selection behavior. + +### Node-image task configuration + +`src/beakerhub/app.py:26-72` defines application traits for node-image import +tasks: + +- `task_backoff_limit` +- `task_active_deadline_seconds` +- `task_ttl_seconds_after_finished` +- `task_node_selector` +- `task_tolerations` +- `task_namespace` + +The reporter image and resource dictionaries can apply to other container +runtimes, but the retry, TTL, node selector, toleration, Pod, Job, and namespace +semantics are Kubernetes-specific. The task subsystem consumes these traits +directly rather than through a runtime task-backend interface. + +`task_namespace` is also used outside the task subsystem by the admin cluster +dashboard and session-log endpoint. It currently acts as the general +Kubernetes namespace despite its task-specific name. + +## Node-image import tasks + +### Kubernetes client construction + +`src/beakerhub/nodes/tasks.py:13-29` imports the Kubernetes Python client, tries +in-cluster configuration, falls back to kubeconfig, and constructs +`BatchV1Api` and `CoreV1Api` clients. + +The module has no runtime-neutral task abstraction. All create, inspect, and +delete operations use Kubernetes clients and resource types. + +### Task creation + +`src/beakerhub/nodes/tasks.py:32-144` implements an import task as a Kubernetes +`batch/v1` Job: + +- An init container runs `beaker context dump` in the selected node image. +- A reporter container reads the result and sends it to the hub callback. +- Both containers exchange output through an `emptyDir` volume. +- The Pod uses Kubernetes resource requirements, restart policy, node selector, + and tolerations. +- The Job uses Kubernetes metadata and labels, a backoff limit, an active + deadline, and a TTL after completion. +- Submission uses `create_namespaced_job`. + +The two-container workflow and callback protocol are conceptually portable, +but their current orchestration and shared-output mechanism are Kubernetes +resource definitions. + +### Task polling and failure diagnosis + +`src/beakerhub/nodes/tasks.py:147-222` polls a Kubernetes Job with +`read_namespaced_job`. It maps Kubernetes Job counters to BeakerHub task states: + +- `status.succeeded` becomes `completed`. +- `status.failed` becomes `failed`. +- `status.active` becomes `running`. +- No counter becomes `pending`. + +Failure diagnosis finds a Pod through the `job-name` label and examines init +container, main container, and Pod status objects. The generated messages use +Kubernetes terms such as init container, container, and Pod phase. + +### Task deletion and resources + +`src/beakerhub/nodes/tasks.py:225-250` deletes a namespaced Job with Kubernetes +background propagation and converts resource dictionaries to +`V1ResourceRequirements`. + +`delete_job` is not currently called elsewhere under `src/beakerhub`, but it is +still part of the task module's backend surface. + +### Import orchestration + +`src/beakerhub/nodes/import_handlers.py` binds the general import workflow to +the Kubernetes task implementation: + +- Line 21 imports `beakerhub.nodes.tasks` as `k8s_tasks`. +- Lines 34-57 describe the operation and errors as Kubernetes Job operations. +- Lines 81-94 read `task_namespace` and call `create_import_job` directly. +- Lines 95-101 persist and expose a `Failed to create K8s Job` error. +- Lines 103-108 store and log the returned identifier as `job_name`. +- Lines 162-172 call `get_job_status` directly while a task is running. +- Line 199 describes the callback caller as a Kubernetes Pod. + +There is also a namespace inconsistency relevant to this coupling. Creation +passes the configured `task_namespace`, but line 165 polls without passing a +namespace. Polling therefore uses the hard-coded `"beakerhub"` default from +`nodes/tasks.py`. + +The callback endpoint and callback-token protocol are otherwise runtime +neutral. They receive JSON from a reporter and update the database without +using a Kubernetes API. + +## Admin cluster dashboard + +`src/beakerhub/dashboard_handlers.py:127-509` implements the cluster dashboard +as a direct Kubernetes cluster inspection endpoint. + +### Client and resource assumptions + +Lines 151-167 load in-cluster configuration or kubeconfig, construct Core and +Batch API clients, and use `task_namespace` as the inspected namespace. + +Lines 169-195 always build the response from these Kubernetes resource groups: + +- Pods +- PersistentVolumeClaims +- Jobs +- Events +- Nodes +- Helm releases stored as Secrets + +If Kubernetes configuration is unavailable, the endpoint returns +`available: false`. It has no dispatch based on the configured runtime backend. + +### Pod inventory + +Lines 198-238 list namespaced Pods and group them by phase and component. The +classification depends on: + +- `app.kubernetes.io/component` +- JupyterHub's `component=singleuser-server` Pod label +- The `session-` Pod-name prefix +- Hub and proxy substrings in Pod names + +An ECS task cannot appear in this inventory without an alternate data source +and response mapping. + +### Storage, Jobs, and Events + +- Lines 240-257 list PVCs and return Kubernetes capacity, phase, and storage + class fields. +- Lines 259-284 list Kubernetes Jobs selected by + `app.kubernetes.io/name=beakerhub` and aggregate their Job status counters. +- Lines 286-311 list namespaced Kubernetes warning Events from the last hour. + +These response sections describe the Kubernetes deployment as well as user +runtimes. A multi-backend dashboard may therefore need to distinguish platform +deployment information from configured runtime information. + +### Cluster nodes and allocated resources + +Lines 313-413 list Kubernetes Nodes and calculate allocated CPU, memory, and +Pod counts by summing resource requests from running and pending Pods across all +namespaces. + +The returned node metadata depends on Kubernetes labels and node status: + +- `node.kubernetes.io/instance-type` +- `beta.kubernetes.io/instance-type` +- `kubernetes.io/os` +- `kubernetes.io/arch` +- kubelet version +- container runtime version +- Kubernetes capacity and allocatable quantities + +Lines 415-468 parse and format Kubernetes CPU and memory quantity syntax. + +### Helm release discovery + +Lines 470-509 list Kubernetes Secrets with `owner=helm`, then infer Helm +release name, status, revision, and update time from Secret labels and metadata. + +This section concerns how BeakerHub itself is deployed rather than the notebook +runtime, but the backend endpoint still assumes access to a Kubernetes cluster. + +## Session log access + +`src/beakerhub/dashboard_handlers.py:512-591` implements session logs as +Kubernetes Pod logs: + +- It loads in-cluster configuration or kubeconfig. +- It uses `task_namespace` as the notebook namespace. +- It locates the session Pod through the + `hub.jupyter.org/servername=` label. +- It reads a named container with `read_namespaced_pod_log`. +- Its response includes `pod_name`. +- Its errors refer to Pods and containers. + +This endpoint is directly coupled to KubeSpawner's Pod labels and Kubernetes +log API. It cannot find or read an ECS-backed session even if the ECS spawner +successfully starts that session. + +## Persistence and exposed terminology + +`src/beakerhub/orm.py:222-240` describes `NodeImageTask` as tracking Kubernetes +Jobs and stores the external task identifier in `job_name`. + +The same field is exposed by several otherwise runtime-neutral surfaces: + +- `src/beakerhub/nodes/import_handlers.py:67,103,133,163-177` +- `src/beakerhub/cli/main.py:88` +- `src/beakerhub/admin_handlers.py:334,401` +- `src/beakerhub/dashboard_handlers.py:84-103` + +These callers do not themselves use Kubernetes APIs. Their coupling is in the +data contract and user-facing terminology. A runtime-neutral task identifier +would need either a schema/API change or a compatibility mapping to the +existing `job_name` field. + +## Kubernetes-named local provisioner + +`src/beakerhub/spawner/provisioner.py:5-13` defines +`KubernetesLocalProvisioner`. It forwards launched kernel stdout and stderr to +the parent process streams so container logs reach the platform log collector. + +The behavior is useful in both Kubernetes and ECS. Only the class name, +docstring, and registered provisioner name are Kubernetes-specific. + +The notebook configuration generated outside `src` selects the literal +`kubernetes-local-provisioner` name. That configuration is excluded from this +document's file inventory because it is under `helm`, but it is the consumer +that makes this name operationally significant. + +## Text-only spawner references + +The following documentation strings name `BeakerKubeSpawner` as the producer +of the secret-policy environment protocol: + +- `src/beakerhub/services/secrets/__init__.py:3-4` +- `src/beakerhub/services/secrets/beakerhub.py:43-58` + +The environment construction now lives in the shared spawner base, so these +references are stale implementation terminology rather than functional +Kubernetes dependencies. + +## Backend boundaries indicated by the current code + +The implicit assumptions group around three runtime operations that are not +provided by the configured spawner alone: + +| Operation | Current implementation | +| --- | --- | +| Node-image task lifecycle | Kubernetes Job create, poll, diagnose, and delete | +| Runtime inventory and health | Kubernetes Pods, Jobs, Nodes, Events, PVCs, and Helm Secrets | +| Session logs | Kubernetes Pod lookup and Pod log API | + +Default backend selection and Kubernetes-specific task settings are configured +on `BeakerHub` itself. Task identifiers and some API response fields also expose +the Kubernetes implementation beyond these operational boundaries. + +## Proposed backend organization + +The backend should be a namespace and configuration profile, not a single +service that performs all backend operations. It can group related +implementations and provide a coherent set of defaults while leaving each +component independently configurable on the application. + +A backend package could contain implementations for: + +- A JupyterHub spawner +- A node-image task runner +- Dashboard data providers or widgets +- A session log provider +- A node launch-spec resolver +- Runtime-specific secret delivery + +For example, a Kubernetes profile could default these components to Kubernetes +implementations, while an ECS profile could default them to ECS and AWS +implementations. Explicit component configuration should be able to override +individual profile defaults. This permits mixed configurations, such as ECS +notebook sessions with Kubernetes import tasks. + +The intended configuration precedence is: + +1. Component's built-in default +2. Backend profile default +3. Explicit component configuration + +The backend profile should select classes and supply default configuration. It +should not own initialized clients, credentials, mutable runtime state, or the +actual component lifecycle. + +### Package organization and many-to-many mapping + +Provider-first package organization becomes ambiguous when implementations are +independently selectable. AWS is both a vendor and a collection of unrelated +services: selecting ECS does not imply Secrets Manager, CloudWatch, S3, or any +particular combination of them. + +A capability-first implementation tree keeps discovery and ownership clearer: + +```text +beakerhub/backends/ +├── contracts/ +│ ├── task_runner.py +│ ├── dashboard.py +│ ├── log_provider.py +│ ├── secret_vault.py +│ └── node_spec.py +├── task_runners/ +│ ├── kubernetes.py +│ └── aws_ecs.py +├── dashboards/ +│ ├── kubernetes.py +│ └── aws_ecs.py +├── log_providers/ +│ ├── kubernetes.py +│ └── aws_cloudwatch.py +├── secret_vaults/ +│ ├── orm.py +│ └── aws_secrets_manager.py +├── secret_delivery/ +│ ├── environment.py +│ ├── kubernetes.py +│ └── aws_ecs.py +├── node_specs/ +│ ├── kubernetes.py +│ └── aws_ecs.py +└── profiles/ + ├── default.py + ├── kubernetes.py + └── aws_ecs.py +``` + +Under this organization, the AWS Secrets Manager implementation belongs at +`backends/secret_vaults/aws_secrets_manager.py`, not under the ECS package. If +the implementation grows, it can become a package at the same location: + +```text +backends/secret_vaults/aws_secrets_manager/ +├── __init__.py +├── vault.py +├── models.py +└── client.py +``` + +The consequence is that code for one vendor appears under several capability +directories. That matches the normal discovery question: which task runners, +secret vaults, or log providers are available? Vendor-oriented inspection can +still use repository search and consistent prefixes such as `aws_`. + +### Sparse profiles and independent overrides + +A backend profile should be an explicit, potentially sparse mapping from +capabilities to default implementations. For example, an ECS profile could +select defaults for session spawning, background tasks, dashboard data, logs, +and workload-spec resolution: + +```python +class AwsEcsBackendProfile(BackendProfile): + spawner_class = BeakerAwsECSSpawner + task_runner_class = AwsEcsTaskRunner + dashboard_provider_class = AwsEcsDashboardProvider + log_provider_class = AwsCloudWatchLogProvider + node_spec_provider_class = AwsEcsNodeSpecProvider +``` + +It should not select AWS Secrets Manager only because it selects ECS. Secret +storage is orthogonal and should remain independently configurable: + +```python +c.BeakerHub.backend_profile = "aws-ecs" +c.BeakerHub.secret_vault_class = ( + "beakerhub.backends.secret_vaults.aws_secrets_manager." + "AwsSecretsManagerVault" +) +``` + +The default profile can supply the current ORM vault when no explicit vault is +configured. Mixed configurations must not require a new named profile for each +combination. In particular, avoid profile proliferation such as +`aws-ecs-with-secrets-manager-and-cloudwatch`. + +The mapping rules should be explicit: + +1. A profile supplies defaults; it does not own its components. +2. A profile can omit capabilities it does not need to configure. +3. Explicit component configuration overrides the profile. +4. Selecting a vendor or runtime does not imply unrelated services. +5. Each implementation lives under the capability it implements. +6. Built-in aliases need only be unique within their capability. +7. Unsupported capabilities remain unset or use an explicit null + implementation; they do not fabricate equivalent data. + +### Discovery and registration + +Initial discovery should be deterministic rather than based on filesystem +scanning or import side effects. Each capability can provide a small built-in +alias registry: + +```python +BUILTIN_TASK_RUNNERS = { + "kubernetes": KubernetesTaskRunner, + "aws-ecs": AwsEcsTaskRunner, +} + +BUILTIN_SECRET_VAULTS = { + "orm": OrmSecretVault, + "aws-secrets-manager": AwsSecretsManagerVault, +} +``` + +Application configuration can accept either a stable alias or an importable +class path: + +```python +c.BeakerHub.task_runner = "aws-ecs" +c.BeakerHub.secret_vault = "aws-secrets-manager" +``` + +If external implementations become necessary, component-specific Python entry +point groups can extend the same registries: + +```text +beakerhub.task_runners +beakerhub.dashboard_providers +beakerhub.secret_vaults +beakerhub.log_providers +``` + +Separate groups avoid global alias collisions and make it possible to list the +available implementations for one capability without importing every backend +component. + +## Additional interchangeable components + +### Secret vault + +Secret storage is the strongest additional candidate for an independently +configured backend service. The current implementation combines secret +metadata, encrypted value storage, CRUD, context enablement, spawn-time +resolution, and runtime delivery. + +Current ORM coupling includes: + +- `src/beakerhub/orm.py:52-87`, where `EncryptedString` encrypts values with a + process-wide Fernet key. +- `src/beakerhub/orm.py:248-292`, where `NodeSecret` stores the secret value, + environment-variable name, description, policy overrides, and node-image + scope. +- `src/beakerhub/orm.py:140-146`, where `beaker_context_secrets` associates + context enablement with numeric `NodeSecret` IDs. +- `src/beakerhub/admin_handlers.py:696-865`, where handlers query and mutate + `NodeSecret` records directly. +- `src/beakerhub/admin_handlers.py:868-955`, where context-specific enablement + is managed through the ORM association table. +- The configured spawner's option processing, which queries these records to + resolve effective values and policy overrides for a launch. + +An AWS Secrets Manager implementation could take several forms: + +- Store values and metadata entirely in AWS Secrets Manager. +- Keep scope, descriptions, policy overrides, and context relationships in the + ORM while storing only secret values externally. +- Store external secret references in the ORM and let the runtime inject values + directly without exposing plaintext to the hub. + +The third form is materially different from replacing `EncryptedString`. +Vault persistence and secret delivery should therefore be separate contracts, +even if a default implementation provides both. + +A vault service would likely own operations such as: + +- List and retrieve secret metadata +- Create, update, and delete secrets +- Resolve stable secret references +- Retrieve values when the selected delivery mechanism requires plaintext +- Report capabilities, such as whether values can be read back through the + administration API + +The existing context associations require stable local identifiers. A fully +external vault still needs either local reference records or a replacement for +the numeric `node_secret_id` relationship. + +### Secret resolution and delivery + +The runtime currently receives plaintext secret values as environment +variables. Additional environment variables identify which values came from +the vault and carry per-secret policy overrides. + +Potential delivery implementations include: + +- Literal environment values +- Kubernetes `secretRef` entries +- ECS `secrets.valueFrom` references +- Mounted secret files +- Node-side lookup through workload identity + +The delivery component must preserve the node-side policy protocol even when +the hub does not retrieve the value. This includes identifying vault-provided +variables and supplying the policy overrides consumed by +`BeakerhubSecretsManager`. + +Secret delivery may live within each runtime backend because its output must +match the target workload specification. It should still consume a +runtime-neutral vault/reference model rather than query `NodeSecret` directly. + +### Node launch-spec resolution + +`NodeImages` currently represents a runnable node as an OCI image assembled +from registry, repository, and tag: + +- `src/beakerhub/orm.py:186-217` +- `src/beakerhub/spawner/base.py:90-144` + +Both notebook spawning and node-image import tasks consume that representation. +For Kubernetes, an image reference is nearly sufficient to build a Pod. ECS +normally also requires a task definition and associated launch configuration. + +ECS resolution may need to: + +- Select an existing task definition +- Register a task-definition revision for the chosen image +- Map CPU and memory settings to valid ECS combinations +- Configure roles, logging, volumes, ports, and health checks +- Apply runtime-specific secret references +- Manage any generated task-definition revisions + +A shared `NodeSpecResolver`, `WorkloadSpecProvider`, or similarly focused +component could translate a `NodeImages` record and launch options into a +backend-specific workload specification. This would avoid duplicating the same +mapping in the spawner and task runner. + +This boundary can also support digest-pinned images, prebuilt task definitions, +private registry authentication, or future artifact representations without +changing the catalog-facing `NodeImages` API immediately. + +### Context catalog sources + +Context discovery already has a useful runtime-neutral interchange format: + +- `src/beakerhub/utils.py:149-156` defines `InterchangeDump`. +- `src/beakerhub/utils.py:579-581` deserializes the format. +- `src/beakerhub/utils.py:584-872` ingests it into the ORM. +- `src/beakerhub/scripts/seed_contexts.py` imports the older JSON format. + +The current primary source executes a node image and inspects its installed +packages. Other potential sources include: + +- A package or organizational context registry +- A static manifest repository +- S3 or another object store +- A Git repository +- OCI image metadata or attestations that do not require executing the image + +A `ContextSource` or `CatalogSource` could produce `InterchangeDump` objects, +while the existing ingestion, deduplication, and curation logic remains shared. +This is a plausible future extension point rather than a requirement for the +first ECS implementation. + +### Task result transport + +The import-task callback protocol is already mostly independent from +Kubernetes. The shared workflow is: + +1. Create a task record and callback credential. +2. Execute `beaker context dump` through the configured task runner. +3. Deliver an interchange dump to BeakerHub. +4. Authenticate and deserialize the result. +5. Ingest the dump and update task state. + +The current implementation delivers the result through an HTTP callback from a +reporter container. Other implementations could use object storage, a queue, +an event, or output retrieved during polling. + +The task runner should own workload submission, status inspection, +cancellation, and backend-specific diagnostics. Result authentication, +deserialization, and ingestion should remain shared. The mechanism that makes +the completed dump available can be independently configurable if a second +transport is needed. + +## Component and capability summary + +The backend-related component inventory is: + +| Component | Responsibility | Configuration relationship | +| --- | --- | --- | +| Spawner | JupyterHub session lifecycle and reachable server address | Selected independently; backend profile supplies a default | +| Task runner | Submit, inspect, cancel, and diagnose background workloads | Selected independently; backend profile supplies a default | +| Dashboard provider/widgets | Runtime and infrastructure inspection | Selected or enabled independently by server configuration | +| Session log provider | Discover and retrieve logs for a session | Selected independently; backend profile supplies a default | +| Node launch-spec resolver | Translate node catalog records into runtime workload definitions | Shared by a backend's spawner and task runner where useful | +| Secret vault | Store secret metadata, references, and optionally values | Independent configured service | +| Secret delivery | Add values or references and policy metadata to workload specifications | Usually runtime-specific; consumes the configured vault | +| Context catalog source | Produce context interchange dumps | Independent optional source | +| Task result transport | Make completed task output available to shared ingestion | Part of the task subsystem and independently replaceable if needed | + +Components should expose capabilities rather than requiring every backend to +implement equivalent Kubernetes concepts. Configuration and API responses need +to distinguish an unsupported capability from an empty result or a temporary +query failure. + +The dashboard should also distinguish deployment information from runtime +information. A BeakerHub deployment can continue to run on Kubernetes while +its notebook sessions and import tasks run on ECS, so Kubernetes deployment +status and ECS runtime status can both be valid and useful at the same time. diff --git a/pyproject.toml b/pyproject.toml index 216c693..cab638b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ artifacts = [ ] [project.entry-points."jupyter_client.kernel_provisioners"] -kubernetes-local-provisioner = "beakerhub.spawner.provisioner:KubernetesLocalProvisioner" +beakerhub-local-provisioner = "beakerhub.services.spawner.provisioner:BeakerhubLocalProvisioner" [tool.hatch.envs.types] extra-dependencies = [ diff --git a/src/beakerhub/admin_handlers.py b/src/beakerhub/admin_handlers.py index 243783b..ed80ecd 100644 --- a/src/beakerhub/admin_handlers.py +++ b/src/beakerhub/admin_handlers.py @@ -447,7 +447,7 @@ async def post(self): import_error = None app = self.settings.get("app") if app: - from beakerhub.nodes.import_handlers import launch_import_task + from beakerhub.tasks.image_import.task import launch_import_task try: import_task = launch_import_task(self.db, app, node) except Exception as e: diff --git a/src/beakerhub/app.py b/src/beakerhub/app.py index 4fb5848..1b1f0f3 100644 --- a/src/beakerhub/app.py +++ b/src/beakerhub/app.py @@ -3,7 +3,7 @@ import traitlets from jupyterhub.app import JupyterHub -from traitlets import Dict, Integer, List, Unicode, default +from traitlets import Dict, Instance, Type, Unicode, default # Import orm tables so that they can be added. import beakerhub.orm # type: ignore @@ -13,8 +13,20 @@ from beakerhub.handlers import get_override_handlers, HierarchicalStaticHandler, VueSPAHandler from beakerhub.api_handlers import handlers as api_handlers from beakerhub.admin_handlers import admin_handlers -from beakerhub.nodes.import_handlers import import_handlers -from beakerhub.dashboard_handlers import dashboard_handlers +from beakerhub.runtimes.base import BaseRuntime, BaseRuntimeBundle +from beakerhub.runtimes.kubernetes import KubernetesRuntime +from beakerhub.services.dashboard.base import BaseDashboardService +from beakerhub.services.dashboard.aws_ecs_dashboard import AwsEcsDashboardService +from beakerhub.services.dashboard.handlers import handlers as dashboard_handlers +from beakerhub.services.dashboard.kubernetes_dashboard import ( + KubernetesDashboardService, +) +from beakerhub.services.task.base import BaseTaskRunnerService +from beakerhub.services.task.aws_ecs_task_runner import AwsEcsTaskRunnerService +from beakerhub.services.task.handlers import handlers as task_handlers +from beakerhub.services.task.kubernetes_task_runner import ( + KubernetesTaskRunnerService, +) PACKAGE_ROOT = Path(__file__).resolve().parent @@ -27,52 +39,45 @@ class BeakerHub(JupyterHub): description = "Beakerhub version of: \n" + str(JupyterHub.description) example = "Beakerhub version of: \n" + str(JupyterHub.examples) - # ---- Task configuration (for node image import jobs) ---- - task_reporter_image = Unicode( - "beakerhub/task-reporter:latest", + runtime_bundle_class = Type( + allow_none=True, + klass=BaseRuntimeBundle, config=True, - help="Full image reference for the task reporter container.", + help="Configures a default suite of runtime based configuration options." ) - task_reporter_pull_policy = Unicode( - "Always", - config=True, - help="Image pull policy for the task reporter container.", - ) - task_reporter_resources = Dict( - config=True, - help="Resource requests/limits for the task reporter container.", + runtime_bundle = Instance( + klass=BaseRuntimeBundle, + allow_none=True, ) - task_node_image_resources = Dict( - config=True, - help="Resource requests/limits for the node image init container in tasks.", - ) - task_backoff_limit = Integer( - 0, + + runtime_class = Type( + klass=BaseRuntime, config=True, - help="Number of retries before marking a task Job as failed.", + help="Configures a default suite of runtime based configuration options." ) - task_active_deadline_seconds = Integer( - 300, - config=True, - help="Maximum time in seconds a task Job can run before being terminated.", + runtime = Instance( + klass=BaseRuntime, + allow_none=True, ) - task_ttl_seconds_after_finished = Integer( - 600, + + task_runner_class = Type( + klass=BaseTaskRunnerService, config=True, - help="Time in seconds to keep completed/failed Jobs before cleanup.", + help="Task-runner service used for background workloads.", ) - task_node_selector = Dict( - config=True, - help="Node selector for task pods.", + task_runner = Instance( + BaseTaskRunnerService, + allow_none=False, ) - task_tolerations = List( + + dashboard_service_class = Type( + klass=BaseDashboardService, config=True, - help="Tolerations for task pods.", + help="Service used to collect runtime dashboard data and session logs.", ) - task_namespace = Unicode( - "beakerhub", - config=True, - help="Kubernetes namespace for task Jobs.", + dashboard_service = Instance( + BaseDashboardService, + allow_none=False, ) enable_idle_session_culling = traitlets.Bool( @@ -127,10 +132,70 @@ def _default_hub_prefix(self): def _default_authenticator_class(self): return CognitoBotoAuthenticator + @default("runtime_bundle") + def _default_runtime_bundle(self): + if self.runtime_bundle_class is not BaseRuntimeBundle: + return self.runtime_bundle_class(parent=self) + else: + return None + + @default("runtime_class") + def _default_runtime_class(self): + if isinstance(self.runtime_bundle, BaseRuntimeBundle): + return self.runtime_bundle.runtime_class + else: + return KubernetesRuntime + + @default("runtime") + def _default_runtime(self): + return self.runtime_class(parent=self) + @default("spawner_class") def _default_spawner_class(self): - from beakerhub.spawner.kubernetes import BeakerKubeSpawner - return BeakerKubeSpawner + if isinstance(self.runtime_bundle, BaseRuntimeBundle): + return self.runtime_bundle.default_spawner_class + else: + from beakerhub.services.spawner.kubernetes_spawner import BeakerKubeSpawner + return BeakerKubeSpawner + + @default("task_runner_class") + def _default_task_runner_class(self): + if isinstance(self.runtime_bundle, BaseRuntimeBundle): + return self.runtime_bundle.default_task_runner_class + else: + return KubernetesTaskRunnerService + + @default("task_runner") + def _default_task_runner(self): + return self.task_runner_class(parent=self) + + @default("dashboard_service_class") + def _default_dashboard_service_class(self): + if isinstance(self.runtime_bundle, BaseRuntimeBundle): + return self.runtime_bundle.default_dashboard_class + return KubernetesDashboardService + + @default("dashboard_service") + def _default_dashboard_service(self): + return self.dashboard_service_class(parent=self) + + def update_config(self, config): + """Refresh services created before JupyterHub loads its config file. + + Child ``Configurable`` instances inherit their parent's config only at + construction time. JupyterHub loads ``beakerhub_config.py`` after the + application is constructed, so an already-created task runner would + otherwise retain its trait defaults. + """ + super().update_config(config) + if self._trait_values.get("runtime_bundle", None): + self.runtime_bundle.update_config(config) + if "runtime" in self._trait_values: + self.runtime.update_config(config) + if "task_runner" in self._trait_values: + self.task_runner.update_config(config) + if "dashboard_service" in self._trait_values: + self.dashboard_service.update_config(config) @default("config_file") def _default_config_file(self): @@ -154,8 +219,17 @@ def _default_db_url(self): def _default_cookie_secret_file(self): return 'beakerhub_cookie_secret' - # Hoist from subclass to provide from new root app - classes = JupyterHub.classes + @default("classes") + def _default_classes(self): + return [ + self.__class__, + BaseTaskRunnerService, + KubernetesTaskRunnerService, + AwsEcsTaskRunnerService, + BaseDashboardService, + KubernetesDashboardService, + AwsEcsDashboardService, + ] def init_db(self): if self.vault_encryption_key: @@ -171,7 +245,13 @@ def init_handlers(self): """Initialize handlers, including custom Vue SPA handlers.""" super().init_handlers() - override_handlers = get_override_handlers(self.base_url, self.beaker_static_path) + api_handlers + admin_handlers + import_handlers + dashboard_handlers + override_handlers = ( + get_override_handlers(self.base_url, self.beaker_static_path) + + api_handlers + + admin_handlers + + task_handlers + + dashboard_handlers + ) overridden_paths = {handler[0] for handler in override_handlers} overridden_paths.add(r'(.*)') # 404NotFound path overridden_paths.update({path for path, *args in self.handlers if path.startswith("/admin")}) @@ -192,13 +272,13 @@ async def init_role_creation(self): later, in `init_role_assignment`, so the ordering is safe. """ if self.enable_idle_session_culling: - from beakerhub.services.idle_culler import SERVICE_ROLE + from beakerhub.services.periodic_tasks.idle_culler import SERVICE_ROLE self.load_roles = [*self.load_roles, SERVICE_ROLE] await super().init_role_creation() def init_services(self): if self.enable_idle_session_culling: - from beakerhub.services.idle_culler import SERVICE_DEFINITION + from beakerhub.services.periodic_tasks.idle_culler import SERVICE_DEFINITION # Copy the definition. Mutating the module-level dict would make # this method unsafe to call more than once. service = dict(SERVICE_DEFINITION) diff --git a/src/beakerhub/cli/main.py b/src/beakerhub/cli/main.py index d6df6ac..2f10926 100644 --- a/src/beakerhub/cli/main.py +++ b/src/beakerhub/cli/main.py @@ -46,7 +46,7 @@ def import_image(image_slug: str, config_file: str, db_url: str | None): """ from beakerhub.app import BeakerHub from beakerhub.orm import NodeImages, NodeImageTask - from beakerhub.nodes.import_handlers import launch_import_task + from beakerhub.tasks.image_import.task import launch_import_task # Initialize app to load config (reporter image, resources, etc.) app = BeakerHub() diff --git a/src/beakerhub/nodes/import_handlers.py b/src/beakerhub/nodes/import_handlers.py deleted file mode 100644 index b05dbbd..0000000 --- a/src/beakerhub/nodes/import_handlers.py +++ /dev/null @@ -1,271 +0,0 @@ -""" -API handlers for node image import tasks. - -Provides endpoints to: -- Trigger a context import from a node image (admin, creates K8s Job) -- Poll import task status (admin) -- Receive results from the task reporter container (internal, token-auth) -""" -import json -import logging -from datetime import datetime, timezone -from secrets import token_urlsafe -from typing import Any, TYPE_CHECKING - -from tornado import web -from jupyterhub.apihandlers import APIHandler -from jupyterhub.scopes import needs_scope -from sqlalchemy.orm import Session - -from beakerhub.orm import NodeImages, NodeImageTask -from beakerhub.nodes import tasks as k8s_tasks -from beakerhub.utils import ( - InterchangeDump, - deserialize_interchange_dump, - ingest_interchange_dump, -) - -if TYPE_CHECKING: - from beakerhub.app import BeakerHub - -log = logging.getLogger(__name__) - - -def launch_import_task( - db: Session, - app: "BeakerHub", - node_image: NodeImages, -) -> NodeImageTask: - """ - Launch a context import K8s Job for the given node image. - - Creates a NodeImageTask record, builds the callback URL, and submits - the K8s Job. Can be called from handlers, CLI, or anywhere with access - to the db session and app instance. - - Args: - db: SQLAlchemy session. - app: The BeakerHub application instance. - node_image: The NodeImages record to import from. - - Returns: - The created NodeImageTask record. - - Raises: - ValueError: If an import is already running for this image. - RuntimeError: If K8s Job creation fails. - """ - # Check for already-running import - existing_task = db.query(NodeImageTask).filter( - NodeImageTask.node_image_id == node_image.id, - NodeImageTask.task_type == "context_import", - NodeImageTask.status.in_(["pending", "running"]), - ).first() - if existing_task: - raise ValueError( - f"Import already in progress for image '{node_image.slug}' " - f"(job: {existing_task.job_name})" - ) - - # Create task record with callback token - callback_token = token_urlsafe(48) - task = NodeImageTask( - node_image_id=node_image.id, - task_type="context_import", - status="pending", - callback_token=callback_token, - ) - db.add(task) - db.commit() - - # Build callback URL using the hub's internal connect URL - hub_url = getattr(app, "hub_connect_url", "http://localhost:8888") - callback_url = f"{hub_url.rstrip('/')}/api/beakerhub/internal/task-callback/{callback_token}" - namespace = getattr(app, "task_namespace", "beakerhub") - - # Create the K8s Job - try: - job_name = k8s_tasks.create_import_job( - app=app, - node_image=node_image, - callback_url=callback_url, - callback_token=callback_token, - namespace=namespace, - ) - except Exception as e: - task.status = "failed" - task.error = f"Failed to create K8s Job: {e}" - task.updated_at = datetime.now(timezone.utc) - db.commit() - log.exception(f"Failed to create import job for {node_image.slug}") - raise RuntimeError(f"Failed to create import job: {e}") from e - - task.job_name = job_name - task.status = "running" - task.updated_at = datetime.now(timezone.utc) - db.commit() - - log.info(f"Launched import job {job_name} for node image {node_image.slug}") - return task - - -class NodeImageImportHandler(APIHandler): - """Trigger a context import for a node image.""" - - @needs_scope('admin:users') - async def post(self, image_id: str): - """Create a K8s Job to import context data from the node image.""" - node = self.db.query(NodeImages).filter(NodeImages.id == int(image_id)).first() - if not node: - raise web.HTTPError(404, f"Node image not found: {image_id}") - - app = self.settings.get("app") - try: - task = launch_import_task(self.db, app, node) - except ValueError as e: - raise web.HTTPError(409, str(e)) - except RuntimeError as e: - raise web.HTTPError(500, str(e)) - - self.set_header("Content-Type", "application/json") - self.write(json.dumps({ - "task_id": task.id, - "job_name": task.job_name, - "status": task.status, - })) - - -class NodeImageImportStatusHandler(APIHandler): - """Poll the status of an import task.""" - - def compute_etag(self) -> None: - return None - - @needs_scope('admin:users') - async def get(self, image_id: str): - """Get the current import status for a node image.""" - node = self.db.query(NodeImages).filter(NodeImages.id == int(image_id)).first() - if not node: - raise web.HTTPError(404, f"Node image not found: {image_id}") - - # Get the most recent task for this image - task = self.db.query(NodeImageTask).filter( - NodeImageTask.node_image_id == node.id, - NodeImageTask.task_type == "context_import", - ).order_by(NodeImageTask.created_at.desc()).first() - - if not task: - self.set_header("Content-Type", "application/json") - self.write(json.dumps({"status": "none", "message": "No import has been run"})) - return - - # If task is still running, poll the K8s Job for updated status - if task.status == "running" and task.job_name: - try: - job_status = k8s_tasks.get_job_status(task.job_name) - if job_status["status"] == "failed": - task.status = "failed" - task.error = job_status["message"] - task.updated_at = datetime.now(timezone.utc) - self.db.commit() - except Exception as e: - log.warning(f"Failed to poll job status for {task.job_name}: {e}") - - result: dict[str, Any] = { - "task_id": task.id, - "status": task.status, - "job_name": task.job_name, - "created_at": task.created_at.isoformat() if task.created_at else None, - "updated_at": task.updated_at.isoformat() if task.updated_at else None, - } - if task.error: - result["error"] = task.error - if task.result: - result["result"] = task.result - - self.set_header("Content-Type", "application/json") - self.write(json.dumps(result)) - - -class TaskCallbackHandler(APIHandler): - """ - Internal callback endpoint for the task reporter container. - - Authenticates via a per-task token (not admin scope). - Receives the interchange dump JSON and ingests it. - """ - - def check_xsrf_cookie(self): - # Internal endpoint called from K8s pod, no XSRF cookie - return - - def get_current_user(self): - # Override auth — this endpoint uses token-based auth, not session auth - return None - - async def post(self, callback_token: str): - """Receive import results from the reporter container.""" - # Validate callback token - task = self.db.query(NodeImageTask).filter( - NodeImageTask.callback_token == callback_token, - NodeImageTask.status == "running", - ).first() - if not task: - raise web.HTTPError(404, "Invalid or expired callback token") - - node_image = task.node_image - - # Parse the request body - try: - body = json.loads(self.request.body) - except (json.JSONDecodeError, TypeError) as e: - task.status = "failed" - task.error = f"Invalid JSON in callback: {e}" - task.updated_at = datetime.now(timezone.utc) - self.db.commit() - raise web.HTTPError(400, f"Invalid JSON: {e}") - - # The output is an array of InterchangeDump objects (one per package) - if not isinstance(body, list): - body = [body] - - # Ingest each package dump - combined_stats: dict[str, int] = {} - try: - for dump in body: - stats = ingest_interchange_dump( - db=self.db, - dump=dump, - node_image=node_image, - enable_contexts=True, - preserve_curated=True, - ) - for key, value in stats.items(): - combined_stats[key] = combined_stats.get(key, 0) + value - except Exception as e: - self.db.rollback() - task.status = "failed" - task.error = f"Ingestion failed: {e}" - task.updated_at = datetime.now(timezone.utc) - self.db.commit() - log.exception(f"Ingestion failed for task {task.id}") - raise web.HTTPError(500, f"Ingestion failed: {e}") - - # Mark task as completed - task.status = "completed" - task.result = combined_stats - task.error = None - task.updated_at = datetime.now(timezone.utc) - self.db.commit() - - log.info(f"Import completed for node image {node_image.slug}: {combined_stats}") - - self.set_header("Content-Type", "application/json") - self.write(json.dumps({"status": "ok", "stats": combined_stats})) - - -import_handlers = [ - (r"/api/beakerhub/admin/node-images/(\d+)/import", NodeImageImportHandler), - (r"/api/beakerhub/admin/node-images/(\d+)/import-status", NodeImageImportStatusHandler), - (r"/api/beakerhub/internal/task-callback/([^/]+)", TaskCallbackHandler), -] diff --git a/src/beakerhub/nodes/tasks.py b/src/beakerhub/nodes/tasks.py deleted file mode 100644 index 48049a7..0000000 --- a/src/beakerhub/nodes/tasks.py +++ /dev/null @@ -1,250 +0,0 @@ -""" -Kubernetes Job management for node image tasks. - -Creates and monitors K8s Jobs that run commands against node images -(e.g., `beaker context dump` for importing context data). -""" -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING -from uuid import uuid4 - -from kubernetes import client as k8s_client -from kubernetes import config as k8s_config - -if TYPE_CHECKING: - from beakerhub.app import BeakerHub - from beakerhub.orm import NodeImages - -log = logging.getLogger(__name__) - - -def _get_k8s_clients(namespace: str) -> tuple[k8s_client.BatchV1Api, k8s_client.CoreV1Api, str]: - """Load in-cluster config and return batch + core API clients.""" - try: - k8s_config.load_incluster_config() - except k8s_config.ConfigException: - k8s_config.load_kube_config() - return k8s_client.BatchV1Api(), k8s_client.CoreV1Api(), namespace - - -def create_import_job( - app: BeakerHub, - node_image: NodeImages, - callback_url: str, - callback_token: str, - namespace: str = "beakerhub", -) -> str: - """ - Create a K8s Job that runs `beaker context dump` on the given node image - and reports results back to the hub via the reporter container. - - Args: - app: The BeakerHub application instance (for task configuration). - node_image: The NodeImages record to import from. - callback_url: The URL the reporter will POST results to. - callback_token: Auth token for the reporter callback. - namespace: K8s namespace to create the Job in. - - Returns: - The Job name. - """ - batch_api, _, _ = _get_k8s_clients(namespace) - - job_name = f"node-import-{node_image.slug}-{uuid4().hex[:8]}" - node_image_ref = node_image.default_img_string - output_volume_name = "task-output" - stdout_path = "/output/stdout" - stderr_path = "/output/stderr" - - # Init container: run the node image to dump context data - init_container = k8s_client.V1Container( - name="context-dump", - image=node_image_ref, - command=["sh", "-c", f"beaker context dump > {stdout_path} 2> {stderr_path}"], - volume_mounts=[ - k8s_client.V1VolumeMount( - name=output_volume_name, - mount_path="/output", - ) - ], - resources=_build_resource_requirements(app.task_node_image_resources), - ) - - # Main container: reporter reads output and POSTs to callback - reporter_container = k8s_client.V1Container( - name="reporter", - image=app.task_reporter_image, - image_pull_policy=app.task_reporter_pull_policy, - env=[ - k8s_client.V1EnvVar(name="CALLBACK_URL", value=callback_url), - k8s_client.V1EnvVar(name="CALLBACK_TOKEN", value=callback_token), - k8s_client.V1EnvVar(name="STDOUT_PATH", value=stdout_path), - k8s_client.V1EnvVar(name="STDERR_PATH", value=stderr_path), - ], - volume_mounts=[ - k8s_client.V1VolumeMount( - name=output_volume_name, - mount_path="/output", - ) - ], - resources=_build_resource_requirements(app.task_reporter_resources), - ) - - # Pod spec - pod_spec = k8s_client.V1PodSpec( - init_containers=[init_container], - containers=[reporter_container], - volumes=[ - k8s_client.V1Volume( - name=output_volume_name, - empty_dir=k8s_client.V1EmptyDirVolumeSource(), - ) - ], - restart_policy="Never", - node_selector=app.task_node_selector or None, - tolerations=[ - k8s_client.V1Toleration(**t) for t in app.task_tolerations - ] if app.task_tolerations else None, - ) - - # Job spec - job = k8s_client.V1Job( - api_version="batch/v1", - kind="Job", - metadata=k8s_client.V1ObjectMeta( - name=job_name, - namespace=namespace, - labels={ - "app.kubernetes.io/name": "beakerhub", - "app.kubernetes.io/component": "node-image-task", - "beakerhub/task-type": "context-import", - "beakerhub/node-image": node_image.slug, - }, - ), - spec=k8s_client.V1JobSpec( - template=k8s_client.V1PodTemplateSpec( - metadata=k8s_client.V1ObjectMeta( - labels={ - "app.kubernetes.io/name": "beakerhub", - "app.kubernetes.io/component": "node-image-task", - }, - ), - spec=pod_spec, - ), - backoff_limit=app.task_backoff_limit, - active_deadline_seconds=app.task_active_deadline_seconds, - ttl_seconds_after_finished=app.task_ttl_seconds_after_finished, - ), - ) - - batch_api.create_namespaced_job(namespace=namespace, body=job) - log.info(f"Created import job {job_name} for node image {node_image.slug}") - return job_name - - -def get_job_status(job_name: str, namespace: str = "beakerhub") -> dict: - """ - Query the status of a K8s Job. - - Returns a dict with: - status: "pending" | "running" | "completed" | "failed" - message: Human-readable status message - """ - batch_api, core_api, _ = _get_k8s_clients(namespace) - - try: - job = batch_api.read_namespaced_job(name=job_name, namespace=namespace) - except k8s_client.ApiException as e: - if e.status == 404: - return {"status": "failed", "message": f"Job {job_name} not found"} - raise - - status = job.status - if status.succeeded and status.succeeded > 0: - return {"status": "completed", "message": "Job completed successfully"} - - if status.failed and status.failed > 0: - message = _get_job_failure_message(core_api, job_name, namespace) - return {"status": "failed", "message": message} - - if status.active and status.active > 0: - return {"status": "running", "message": "Job is running"} - - return {"status": "pending", "message": "Job is pending"} - - -def _get_job_failure_message( - core_api: k8s_client.CoreV1Api, - job_name: str, - namespace: str, -) -> str: - """Extract a useful error message from a failed Job's Pod events/status.""" - try: - pods = core_api.list_namespaced_pod( - namespace=namespace, - label_selector=f"job-name={job_name}", - ) - except k8s_client.ApiException: - return "Job failed (could not retrieve pod details)" - - if not pods.items: - return "Job failed (no pods found)" - - pod = pods.items[0] - - # Check init container statuses first (where context-dump runs) - for cs in (pod.status.init_container_statuses or []): - if cs.state and cs.state.waiting: - reason = cs.state.waiting.reason or "Unknown" - msg = cs.state.waiting.message or "" - return f"Init container '{cs.name}' waiting: {reason}. {msg}".strip() - if cs.state and cs.state.terminated and cs.state.terminated.exit_code != 0: - reason = cs.state.terminated.reason or "Error" - msg = cs.state.terminated.message or "" - return f"Init container '{cs.name}' failed ({reason}, exit code {cs.state.terminated.exit_code}). {msg}".strip() - - # Check main container statuses - for cs in (pod.status.container_statuses or []): - if cs.state and cs.state.waiting: - reason = cs.state.waiting.reason or "Unknown" - msg = cs.state.waiting.message or "" - return f"Container '{cs.name}' waiting: {reason}. {msg}".strip() - if cs.state and cs.state.terminated and cs.state.terminated.exit_code != 0: - reason = cs.state.terminated.reason or "Error" - msg = cs.state.terminated.message or "" - return f"Container '{cs.name}' failed ({reason}, exit code {cs.state.terminated.exit_code}). {msg}".strip() - - # Fallback: check pod-level conditions - phase = pod.status.phase or "Unknown" - reason = pod.status.reason or "" - return f"Job failed (pod phase: {phase}). {reason}".strip() - - -def delete_job(job_name: str, namespace: str = "beakerhub") -> None: - """Delete a Job and its associated Pods.""" - batch_api, _, _ = _get_k8s_clients(namespace) - try: - batch_api.delete_namespaced_job( - name=job_name, - namespace=namespace, - body=k8s_client.V1DeleteOptions( - propagation_policy="Background", - ), - ) - except k8s_client.ApiException as e: - if e.status != 404: - raise - - -def _build_resource_requirements( - resources: dict, -) -> k8s_client.V1ResourceRequirements | None: - """Convert a resource dict from config into a K8s ResourceRequirements object.""" - if not resources: - return None - return k8s_client.V1ResourceRequirements( - requests=resources.get("requests"), - limits=resources.get("limits"), - ) diff --git a/src/beakerhub/orm.py b/src/beakerhub/orm.py index 96d2220..d0d27d6 100644 --- a/src/beakerhub/orm.py +++ b/src/beakerhub/orm.py @@ -217,6 +217,27 @@ def default_img_string(self) -> str: return f"{self.default_registry}/{self.repository}:{self.default_tag}" +class BeakerTask(Base): + """Tracks an asynchronous BeakerHub task across runner interactions. + + ``external_id`` is the runtime identifier assigned by the configured task + runner, such as a Kubernetes Job name or ECS task ARN. It is null before + submission succeeds. + """ + __tablename__ = 'beaker_tasks' + id = Column(Integer, primary_key=True) + + external_id = Column(Unicode(255), nullable=True, index=True) + task_type = Column(Unicode(64), nullable=False, index=True) + task_definition = Column(JSONDict, nullable=True) + status = Column(Unicode(32), nullable=False, default="pending") + result = Column(JSONDict, nullable=True) + error = Column(Unicode(4096), nullable=True) + + created_at = Column(DateTime, default=utcnow) + updated_at = Column(DateTime, default=utcnow) + + class NodeImageTask(Base): """ Tracks tasks (K8s Jobs) run against node images. @@ -237,7 +258,6 @@ class NodeImageTask(Base): task_type = Column(Unicode(64), nullable=False) # "context_import", etc. status = Column(Unicode(32), nullable=False, default="pending") # pending, running, completed, failed job_name = Column(Unicode(255), nullable=True) # K8s Job name - callback_token = Column(Unicode(128), nullable=True) # Auth token for reporter callback result = Column(JSONDict, nullable=True) # Ingestion stats on success error = Column(Unicode(4096), nullable=True) # Error message on failure diff --git a/src/beakerhub/runtimes/__init__.py b/src/beakerhub/runtimes/__init__.py new file mode 100644 index 0000000..a4e0706 --- /dev/null +++ b/src/beakerhub/runtimes/__init__.py @@ -0,0 +1,21 @@ +"""Execution-provider runtime abstractions for BeakerHub.""" + +from .base import ( + BaseDefinition, + BaseProcess, + BaseRuntime, + BaseRuntimeBundle, + ProcessOutput, + ProcessStatus, + ProcessType, +) + +__all__ = [ + "BaseDefinition", + "BaseProcess", + "BaseRuntime", + "BaseRuntimeBundle", + "ProcessOutput", + "ProcessStatus", + "ProcessType", +] diff --git a/src/beakerhub/runtimes/aws_ecs.py b/src/beakerhub/runtimes/aws_ecs.py new file mode 100644 index 0000000..f456e92 --- /dev/null +++ b/src/beakerhub/runtimes/aws_ecs.py @@ -0,0 +1,731 @@ +"""AWS ECS runtime workloads. + +This module owns ECS control-plane mechanics. Task-runner and JupyterHub +spawner adapters should retain their respective persistence and lifecycle +policies while delegating provider operations to these classes. +""" + +import hashlib +import re +from typing import Any, Literal + +import boto3 +import traitlets +from botocore.client import BaseClient +from botocore.exceptions import ClientError + +from beakerhub.runtimes.base import ( + BaseDefinition, + BaseProcess, + BaseRuntime, + BaseRuntimeBundle, + ProcessOutput, + ProcessStatus, + ProcessType, +) + + +class AwsEcsRuntime(BaseRuntime): + """Shared ECS and CloudWatch clients for ECS-backed workloads.""" + + ecs_client: BaseClient = traitlets.Instance(klass=BaseClient, config=False) + logs_client: BaseClient = traitlets.Instance(klass=BaseClient, config=False) + aws_region: str = traitlets.Unicode( + help="AWS Region in which the ECS cluster resides.", + config=True, + ) + + @traitlets.default("ecs_client") + def _default_ecs_client(self) -> BaseClient: + return boto3.client("ecs", region_name=self.aws_region or None) + + @traitlets.default("logs_client") + def _default_logs_client(self) -> BaseClient: + return boto3.client("logs", region_name=self.aws_region or None) + + @traitlets.default("aws_region") + def _default_aws_region(self) -> str: + session = boto3._get_default_session() + return (session.region_name if session else None) or "us-east-1" + + # These provider-wide values are copied into a process only when that + # process does not receive an explicit override. This keeps one runtime + # bundle from repeating the same deployment configuration per service. + log_group: str = traitlets.Unicode( + default_value="", + help="Default CloudWatch Logs group for ECS containers.", + config=True, + ) + log_stream_prefix: str = traitlets.Unicode( + default_value="beakerhub-task", + help="Default CloudWatch Logs stream prefix for ECS containers.", + config=True, + ) + execution_role_arn: str = traitlets.Unicode( + default_value="", + help="Default IAM execution role for dynamically registered definitions.", + config=True, + ) + task_role_arn: str = traitlets.Unicode( + default_value="", + help="Default IAM role for dynamically registered task containers.", + config=True, + ) + subnets: list[str] = traitlets.List( + traitlets.Unicode(), + default_value=[], + help="Default subnets for Fargate awsvpc workloads.", + config=True, + ) + security_groups: list[str] = traitlets.List( + traitlets.Unicode(), + default_value=[], + help="Default security groups for Fargate awsvpc workloads.", + config=True, + ) + assign_public_ip: bool = traitlets.Bool( + False, + help="Default public-IP assignment for Fargate workloads.", + config=True, + ) + cpu_architecture: Literal["X86_64", "ARM64"] = traitlets.Enum( + values=["X86_64", "ARM64"], + default_value="X86_64", + help="Default CPU architecture for dynamically registered definitions.", + config=True, + ) + task_definition_name: str = traitlets.Unicode( + default_value="beakerhub-tasks", + help="Default family for dynamically registered ECS task definitions.", + config=True, + ) + cluster_name: str = traitlets.Unicode( + help="Default ECS cluster name or ARN.", + config=True, + ) + task_group: str = traitlets.Unicode( + help="Default ECS task group.", + config=True, + ) + launch_type: Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"] = ( + traitlets.Enum( + values=["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], + default_value="FARGATE", + help="Default ECS launch type.", + config=True, + ) + ) + launch_options: dict[str, Any] = traitlets.Dict( + default_value={}, + help="Default ECS run_task options.", + config=True, + ) + task_overrides: dict[str, Any] | None = traitlets.Dict( + default_value=None, + allow_none=True, + help="Default ECS task overrides.", + config=True, + ) + default_task_cpu: str = traitlets.Unicode( + default_value="4 vcpu", + help="Default CPU for dynamically registered task definitions.", + config=True, + ) + default_task_memory: str = traitlets.Unicode( + default_value="16GB", + help="Default memory for dynamically registered task definitions.", + config=True, + ) + default_task_ephemeral_storage: int = traitlets.Int( + default_value=50, + help="Default ephemeral storage in GiB for Fargate definitions.", + config=True, + ) + + def run_task(self, request: dict[str, Any]) -> str: + """Run one ECS task and return its ARN.""" + try: + response = self.ecs_client.run_task(**request) + except ClientError as error: + raise RuntimeError(f"ECS failed to run task: {error}") from error + tasks = response.get("tasks", []) + if not tasks: + raise RuntimeError(f"ECS did not start a task: {response.get('failures', [])!r}") + return tasks[0]["taskArn"] + + def describe_task(self, cluster_name: str, external_id: str) -> dict[str, Any] | None: + """Return an ECS task description, or ``None`` when it no longer exists.""" + try: + response = self.ecs_client.describe_tasks( + cluster=cluster_name, + tasks=[external_id], + ) + except ClientError as error: + raise RuntimeError(f"ECS failed to describe task {external_id}: {error}") from error + tasks = response.get("tasks", []) + return tasks[0] if tasks else None + + def stop_task(self, cluster_name: str, external_id: str, reason: str) -> None: + """Request that ECS stop a task.""" + try: + self.ecs_client.stop_task( + cluster=cluster_name, + task=external_id, + reason=reason, + ) + except ClientError as error: + raise RuntimeError(f"ECS failed to stop task {external_id}: {error}") from error + + @staticmethod + def awsvpc_network_configuration( + subnets: list[str], + security_groups: list[str], + assign_public_ip: bool, + ) -> dict[str, Any]: + """Build an ECS awsvpc network configuration.""" + return { + "awsvpcConfiguration": { + "subnets": list(subnets), + "securityGroups": list(security_groups), + "assignPublicIp": "ENABLED" if assign_public_ip else "DISABLED", + } + } + + def get_log_events( + self, + log_group: str, + stream_name: str, + ) -> list[str]: + """Read all currently retained events from a CloudWatch log stream.""" + messages: list[str] = [] + next_token: str | None = None + while True: + request: dict[str, Any] = { + "logGroupName": log_group, + "logStreamName": stream_name, + "startFromHead": True, + } + if next_token: + request["nextToken"] = next_token + try: + response = self.logs_client.get_log_events(**request) + except ClientError as error: + raise RuntimeError( + f"CloudWatch failed to get output from {stream_name}: {error}" + ) from error + messages.extend(event["message"] for event in response.get("events", [])) + token = response.get("nextForwardToken") + if not token or token == next_token: + return messages + next_token = token + + def get_log_tail( + self, + log_group: str, + stream_name: str, + tail_lines: int, + ) -> tuple[list[str], bool]: + """Read recent events and report whether older events remain.""" + messages: list[str] = [] + next_token: str | None = None + while True: + request: dict[str, Any] = { + "logGroupName": log_group, + "logStreamName": stream_name, + "startFromHead": False, + "limit": min(10_000, tail_lines - len(messages)), + } + if next_token: + request["nextToken"] = next_token + try: + response = self.logs_client.get_log_events(**request) + except ClientError as error: + raise RuntimeError( + f"CloudWatch failed to get output from {stream_name}: {error}" + ) from error + messages[0:0] = [ + event["message"] for event in response.get("events", []) + ] + token = response.get("nextBackwardToken") + has_older_events = bool(token and token != next_token) + if len(messages) >= tail_lines or not has_older_events: + return messages[-tail_lines:], has_older_events + next_token = token + + +class AwsEcsDefinition(BaseDefinition): + """An ECS workload definition. + + Set ``task_definition_arn`` for a deployment-managed ECS definition, as session + spawners normally do. Without it, :class:`AwsEcsProcess` registers a + definition from the image fields before launch, as task runners need. + """ + + runtime = traitlets.Instance(AwsEcsRuntime, allow_none=True) + task_definition_arn = traitlets.Unicode(allow_none=True, default_value=None, config=True) + container_name = traitlets.Unicode("task-container", config=True) + cpu = traitlets.Unicode(config=True) + memory = traitlets.Unicode(config=True) + ephemeral_storage = traitlets.Int(config=True) + tags = traitlets.Dict(traitlets.Unicode(), traitlets.Unicode(), default_value={}, config=True) + definition_tags = traitlets.Dict( + traitlets.Unicode(), traitlets.Unicode(), default_value={}, config=True + ) + log_group = traitlets.Unicode(config=True) + log_stream_prefix = traitlets.Unicode(config=True) + execution_role_arn = traitlets.Unicode(config=True) + task_role_arn = traitlets.Unicode(config=True) + subnets = traitlets.List(traitlets.Unicode(), config=True) + security_groups = traitlets.List(traitlets.Unicode(), config=True) + assign_public_ip = traitlets.Bool(config=True) + cpu_architecture = traitlets.Enum(["X86_64", "ARM64"], config=True) + task_definition_name = traitlets.Unicode(config=True) + cluster_name = traitlets.Unicode(config=True) + task_group = traitlets.Unicode(config=True) + launch_type = traitlets.Enum( + ["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], config=True + ) + launch_options = traitlets.Dict(config=True) + task_overrides = traitlets.Dict(allow_none=True, config=True) + volumes = traitlets.List(traitlets.Dict(), default_value=[], config=True) + mount_points = traitlets.List(traitlets.Dict(), default_value=[], config=True) + sidecar_containers = traitlets.List(traitlets.Dict(), default_value=[], config=True) + + @traitlets.default("cpu") + def _default_cpu(self) -> str: + return self.runtime.default_task_cpu + + @traitlets.default("memory") + def _default_memory(self) -> str: + return self.runtime.default_task_memory + + @traitlets.default("ephemeral_storage") + def _default_ephemeral_storage(self) -> int: + return self.runtime.default_task_ephemeral_storage + + @traitlets.default("log_group") + def _default_log_group(self) -> str: + return self.runtime.log_group + + @traitlets.default("log_stream_prefix") + def _default_log_stream_prefix(self) -> str: + return self.runtime.log_stream_prefix + + @traitlets.default("execution_role_arn") + def _default_execution_role_arn(self) -> str: + return self.runtime.execution_role_arn + + @traitlets.default("task_role_arn") + def _default_task_role_arn(self) -> str: + return self.runtime.task_role_arn + + @traitlets.default("subnets") + def _default_subnets(self) -> list[str]: + return list(self.runtime.subnets) + + @traitlets.default("security_groups") + def _default_security_groups(self) -> list[str]: + return list(self.runtime.security_groups) + + @traitlets.default("assign_public_ip") + def _default_assign_public_ip(self) -> bool: + return self.runtime.assign_public_ip + + @traitlets.default("cpu_architecture") + def _default_cpu_architecture(self) -> Literal["X86_64", "ARM64"]: + return self.runtime.cpu_architecture + + @traitlets.default("task_definition_name") + def _default_task_definition_name(self) -> str: + base = re.sub(r"[^A-Za-z0-9_-]", "_", self.runtime.task_definition_name) + image = self.image.rsplit("/", maxsplit=1)[-1] + image_name = re.sub(r"[^A-Za-z0-9_-]", "_", image) + image_hash = hashlib.sha256(self.image.encode()).hexdigest()[:12] + base = base or "beakerhub" + image_name = image_name or "image" + max_base_length = 255 - len(image_hash) - 3 + base = base[:max_base_length] + max_image_length = 255 - len(base) - len(image_hash) - 2 + image_name = image_name[:max_image_length] + return f"{base}_{image_name}_{image_hash}" + + @traitlets.default("cluster_name") + def _default_cluster_name(self) -> str: + return self.runtime.cluster_name + + @traitlets.default("task_group") + def _default_task_group(self) -> str: + return self.runtime.task_group + + @traitlets.default("launch_type") + def _default_launch_type( + self, + ) -> Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"]: + return self.runtime.launch_type + + @traitlets.default("launch_options") + def _default_launch_options(self) -> dict[str, Any]: + return dict(self.runtime.launch_options) + + @traitlets.default("task_overrides") + def _default_task_overrides(self) -> dict[str, Any] | None: + return ( + dict(self.runtime.task_overrides) + if self.runtime.task_overrides is not None + else None + ) + + def validate_configuration(self) -> None: + """Validate this definition's ECS launch configuration.""" + if not self.image: + raise ValueError("AwsEcsDefinition.image must be configured") + if not self.cluster_name.strip(): + raise ValueError("AwsEcsDefinition.cluster_name must be configured") + if self.launch_type == "FARGATE": + network_configuration = self.launch_options.get("networkConfiguration") + subnets = ( + network_configuration.get("awsvpcConfiguration", {}).get("subnets", []) + if network_configuration + else self.subnets + ) + if not subnets: + raise ValueError( + "AwsEcsDefinition.subnets must be configured for Fargate workloads" + ) + elif self.ephemeral_storage and not self.task_definition_arn: + raise ValueError( + "ephemeral_storage is supported only for Fargate workloads" + ) + if ( + self.launch_type == "MANAGED_INSTANCES" + and not self.launch_options.get("capacityProviderStrategy") + ): + raise ValueError( + "Managed Instances workloads require launch_options.capacityProviderStrategy" + ) + + def find_or_register_task_definition(self) -> str: + """Return an equivalent ECS task-definition ARN, registering it if needed.""" + if self.task_definition_arn: + return self.task_definition_arn + + expected = self._normalize_task_definition(self.aws_task_definition) + next_token: str | None = None + while True: + request: dict[str, Any] = { + "familyPrefix": self.task_definition_name, + "sort": "DESC", + } + if next_token: + request["nextToken"] = next_token + response = self.runtime.ecs_client.list_task_definitions(**request) + for arn in response.get("taskDefinitionArns", []): + described = self.runtime.ecs_client.describe_task_definition( + taskDefinition=arn + ) + existing = described.get("taskDefinition", described) + if self._normalize_task_definition(existing) == expected: + self.task_definition_arn = existing.get("taskDefinitionArn", arn) + return self.task_definition_arn + next_token = response.get("nextToken") + if not next_token: + break + + try: + response = self.runtime.ecs_client.register_task_definition( + **self.aws_task_definition + ) + except ClientError as error: + raise RuntimeError(f"ECS failed to register task definition: {error}") from error + self.task_definition_arn = response["taskDefinition"]["taskDefinitionArn"] + return self.task_definition_arn + + @staticmethod + def _normalize_task_definition(definition: dict[str, Any]) -> dict[str, Any]: + """Remove ECS-generated fields and empty defaults before comparison.""" + generated_fields = { + "taskDefinitionArn", + "revision", + "status", + "requiresAttributes", + "compatibilities", + "registeredAt", + "registeredBy", + "deregisteredAt", + "tags", + } + + def normalize(value: Any) -> Any: + if isinstance(value, dict): + return { + key: normalized + for key, item in value.items() + if key not in generated_fields + and (normalized := normalize(item)) not in (None, [], {}, False) + } + if isinstance(value, list): + return [normalize(item) for item in value] + return value + + return normalize(definition) + + @property + def aws_task_definition(self): + container: dict[str, Any] = { + "name": self.container_name, + "image": self.image, + "essential": True, + } + if self.entrypoint: + container["entryPoint"] = list(self.entrypoint) + if self.command: + container["command"] = list(self.command) + if self.working_directory: + container["workingDirectory"] = self.working_directory + if self.environment: + container["environment"] = [ + {"name": name, "value": value} + for name, value in self.environment.items() + ] + if self.mount_points: + container["mountPoints"] = [dict(mount) for mount in self.mount_points] + if self.sidecar_containers: + container["dependsOn"] = [ + {"containerName": sidecar["name"], "condition": "SUCCESS"} + for sidecar in self.sidecar_containers + ] + if self.log_group: + container["logConfiguration"] = { + "logDriver": "awslogs", + "options": { + "awslogs-group": self.log_group, + "awslogs-region": self.runtime.aws_region, + "awslogs-stream-prefix": self.log_stream_prefix, + }, + } + + definition_dict: dict[str, Any] = { + "family": self.task_definition_name, + "containerDefinitions": [container, *self.sidecar_containers], + "cpu": self.cpu, + "memory": self.memory, + "runtimePlatform": { + "cpuArchitecture": self.cpu_architecture, + "operatingSystemFamily": "LINUX", + }, + } + if self.volumes: + definition_dict["volumes"] = [dict(volume) for volume in self.volumes] + if self.launch_type == "FARGATE": + definition_dict.update( + { + "ephemeralStorage": { + "sizeInGiB": self.ephemeral_storage + }, + "networkMode": "awsvpc", + "requiresCompatibilities": ["FARGATE"], + } + ) + if self.definition_tags: + definition_dict["tags"] = [ + {"key": key, "value": value} + for key, value in self.definition_tags.items() + ] + if self.execution_role_arn: + definition_dict["executionRoleArn"] = self.execution_role_arn + if self.task_role_arn: + definition_dict["taskRoleArn"] = self.task_role_arn + return definition_dict + +class AwsEcsProcess(BaseProcess): + """A single launched ECS workload.""" + + definition: AwsEcsDefinition + runtime = traitlets.Instance(AwsEcsRuntime, allow_none=False) + + def __init__(self, definition: AwsEcsDefinition, **kwargs: Any) -> None: + if not isinstance(definition, AwsEcsDefinition): + raise TypeError("AwsEcsProcess requires an AwsEcsDefinition") + super().__init__(definition, **kwargs) + + @classmethod + def start( + cls, + definition: AwsEcsDefinition, + *, + process_type: ProcessType = "task", + **kwargs: Any, + ) -> "AwsEcsProcess": + """Launch an ECS workload from a fixed or dynamically registered definition.""" + self = cls(definition, process_type=process_type, **kwargs) + definition.validate_configuration() + task_definition_arn = definition.find_or_register_task_definition() + external_id = self.runtime.run_task( + self._run_task_request(task_definition_arn) + ) + self.external_id = external_id + return self + + def _run_task_request(self, task_definition_arn: str) -> dict[str, Any]: + definition = self.definition + tags = dict(definition.tags) + tags.setdefault("beakerhub-process", "true") + tags.setdefault("beakerhub-process-type", self.process_type) + request: dict[str, Any] = { + "cluster": definition.cluster_name, + "taskDefinition": task_definition_arn, + "overrides": self._build_overrides(), + "tags": [{"key": key, "value": value} for key, value in tags.items()], + } + if definition.task_group: + request["group"] = definition.task_group + if ( + definition.launch_type == "FARGATE" + and "networkConfiguration" not in definition.launch_options + ): + request["networkConfiguration"] = self.runtime.awsvpc_network_configuration( + definition.subnets, + definition.security_groups, + definition.assign_public_ip, + ) + if "capacityProviderStrategy" not in definition.launch_options: + request["launchType"] = definition.launch_type + request.update(definition.launch_options) + return request + + def _build_overrides(self) -> dict[str, Any]: + definition = self.definition + overrides = dict(definition.task_overrides or {}) + containers: list[dict[str, Any]] = [] + task_container: dict[str, Any] = {"name": definition.container_name} + environment: dict[str, str] = {} + for container in overrides.pop("containerOverrides", []): + container = dict(container) + if container.get("name") != definition.container_name: + containers.append(container) + continue + environment.update( + {item["name"]: item["value"] for item in container.pop("environment", [])} + ) + task_container.update(container) + environment.update(definition.environment) + if environment: + task_container["environment"] = [ + {"name": name, "value": value} for name, value in environment.items() + ] + containers.append(task_container) + overrides["containerOverrides"] = containers + return overrides + + def describe(self) -> ProcessStatus: + """Return normalized lifecycle status for this ECS workload.""" + if not self.external_id: + return ProcessStatus("pending", "ECS task has not been submitted") + task = self.runtime.describe_task(self.definition.cluster_name, self.external_id) + if task is None: + return ProcessStatus("failed", f"ECS task {self.external_id} was not found") + status = task.get("lastStatus", "UNKNOWN") + if status in {"PROVISIONING", "PENDING", "ACTIVATING"}: + return ProcessStatus("pending", f"ECS task is {status.lower()}") + if status in {"RUNNING", "DEACTIVATING", "DEPROVISIONING", "STOPPING"}: + return ProcessStatus("running", f"ECS task is {status.lower()}") + if status != "STOPPED": + return ProcessStatus("failed", f"ECS task has unexpected status {status!r}") + + container = self._task_container(task) + exit_code = (container or {}).get("exitCode") + if exit_code == 0: + return ProcessStatus( + "completed", + "ECS task completed successfully", + exit_code=exit_code, + ) + reason = ( + (container or {}).get("reason") + or task.get("stoppedReason") + or "Unknown failure" + ) + return ProcessStatus( + "failed", + f"ECS task failed ({reason})", + exit_code=exit_code, + ) + + def collect_output(self) -> ProcessOutput | None: + """Return retained CloudWatch output for this ECS workload.""" + if not self.definition.log_group or not self.external_id: + return None + task = self.runtime.describe_task(self.definition.cluster_name, self.external_id) + if task is None: + return None + container = self._task_container(task) + task_id = self.external_id.rsplit("/", maxsplit=1)[-1] + container_name = (container or {}).get("name", self.definition.container_name) + stream_name = (container or {}).get( + "logStreamName", + f"{self.definition.log_stream_prefix}/{container_name}/{task_id}", + ) + return ProcessOutput( + stdout="\n".join(self.runtime.get_log_events(self.definition.log_group, stream_name)), + stderr="", + ) + + def stop(self) -> None: + """Request normal cleanup of this ECS workload.""" + if not self.external_id: + return + self.runtime.stop_task( + self.definition.cluster_name, + self.external_id, + f"Stopped BeakerHub {self.process_type} cleanup", + ) + + def _task_container(self, task: dict[str, Any]) -> dict[str, Any] | None: + return next( + ( + item + for item in task.get("containers", []) + if item.get("name") == self.definition.container_name + ), + None, + ) + + +class AwsEcsRuntimeBundle(BaseRuntimeBundle): + """Composition bundle for ECS workloads.""" + + runtime_class = traitlets.Type( + klass=AwsEcsRuntime, + default_value=AwsEcsRuntime, + config=True, + ) + process_class = traitlets.Type( + klass=AwsEcsProcess, + default_value=AwsEcsProcess, + config=True, + ) + definition_class = traitlets.Type( + klass=AwsEcsDefinition, + default_value=AwsEcsDefinition, + config=True, + ) + + default_dashboard_class = traitlets.Type( + klass="beakerhub.services.dashboard.aws_ecs_dashboard.AwsEcsDashboardService", + default_value="beakerhub.services.dashboard.aws_ecs_dashboard.AwsEcsDashboardService", + config=True + ) + default_spawner_class = traitlets.Type( + klass="beakerhub.services.spawner.aws_ecs_spawner.BeakerAwsECSSpawner", + default_value="beakerhub.services.spawner.aws_ecs_spawner.BeakerAwsECSSpawner", + config=True + ) + default_task_runner_class = traitlets.Type( + klass="beakerhub.services.task.aws_ecs_task_runner.AwsEcsTaskRunnerService", + default_value="beakerhub.services.task.aws_ecs_task_runner.AwsEcsTaskRunnerService", + config=True + ) + + +# Task persistence/completion handling and JupyterHub proxy/session behavior +# deliberately remain in their respective service adapters. diff --git a/src/beakerhub/runtimes/base.py b/src/beakerhub/runtimes/base.py new file mode 100644 index 0000000..1365620 --- /dev/null +++ b/src/beakerhub/runtimes/base.py @@ -0,0 +1,227 @@ +"""Provider-neutral runtime workload contracts and composition helpers.""" + +import asyncio +import uuid +from dataclasses import dataclass +from time import monotonic +from typing import Any, Literal, TypeAlias + +from traitlets import Dict, Instance, List, Type, Unicode, default +from traitlets.config import LoggingConfigurable + + +ProcessState: TypeAlias = Literal["pending", "running", "completed", "failed"] +ProcessType: TypeAlias = Literal["task", "service"] + + +@dataclass(frozen=True) +class ProcessStatus: + """The runtime's current view of a launched process.""" + + state: ProcessState + message: str | None = None + exit_code: int | None = None + + @property + def done(self) -> bool: + """Return true when the process has reached a terminal state.""" + return self.state in {"completed", "failed"} + + +@dataclass(frozen=True) +class ProcessOutput: + """Captured standard streams for a process.""" + + stdout: str + stderr: str + + +class BaseRuntime(LoggingConfigurable): + """Own a provider connection and provider-wide runtime defaults.""" + + +class BaseDefinition(LoggingConfigurable): + """Provider-neutral description of an OCI workload. + + A definition describes the workload to launch, not its runtime identity or + lifecycle. Its unset provider fields resolve from ``runtime``. Provider + definitions may add scheduling and launch fields. + """ + + runtime = Instance(BaseRuntime, allow_none=True) + image = Unicode(default_value=None, config=True) + entrypoint = List(Unicode(), default_value=[], config=True) + command = List(Unicode(), default_value=[], config=True) + working_directory = Unicode(allow_none=True, default_value=None, config=True) + environment = Dict(Unicode(), Unicode(), default_value={}, config=True) + labels = Dict(Unicode(), Unicode(), default_value={}, config=True) + + + def __init__(self, **kwargs: Any) -> None: + runtime = kwargs.get("runtime") + parent = kwargs.get("parent") + if runtime is None and isinstance(parent, BaseRuntime): + kwargs["runtime"] = parent + elif runtime is not None and parent is None: + kwargs["parent"] = runtime + elif runtime is not None and parent is not runtime: + raise ValueError("A definition's parent and runtime must be the same") + super().__init__(**kwargs) + + +class BaseProcess(LoggingConfigurable): + """A launched provider workload with an external lifecycle identifier.""" + + runtime = Instance(BaseRuntime, allow_none=False) + + def __init__( + self, + definition: BaseDefinition, + *, + external_id: str | None = None, + process_type: ProcessType = "task", + **kwargs: Any, + ) -> None: + runtime = kwargs.get("runtime") + parent = kwargs.get("parent") + if runtime is None and isinstance(parent, BaseRuntime): + kwargs["runtime"] = parent + elif runtime is not None and parent is None: + kwargs["parent"] = runtime + elif runtime is not None and parent is not runtime: + raise ValueError("A process's parent and runtime must be the same") + + self.id = kwargs.pop("id", uuid.uuid4().hex) + self.process_type = process_type + self.external_id = external_id + self.definition = definition + super().__init__(**kwargs) + if definition.runtime is None: + definition.runtime = self.runtime + elif definition.runtime is not self.runtime: + raise ValueError("A process and its definition must use the same runtime") + if definition.parent is None: + definition.parent = self.runtime + elif definition.parent is not self.runtime: + raise ValueError("A definition's parent and runtime must be the same") + + @classmethod + def start( + cls, + definition: BaseDefinition, + *, + process_type: ProcessType = "task", + **kwargs: Any, + ) -> "BaseProcess": + """Launch a process for ``definition`` and return its runtime handle.""" + raise NotImplementedError + + def describe(self) -> ProcessStatus: + """Return the latest provider status for this process.""" + raise NotImplementedError + + @property + def status(self) -> ProcessStatus: + """Return the latest provider status for this process.""" + return self.describe() + + def collect_output(self) -> ProcessOutput | None: + """Return available process output, if the provider retains it.""" + raise NotImplementedError + + def stop(self) -> None: + """Request that the provider stop this process.""" + raise NotImplementedError + + async def await_completion(self, timeout: float | None = 600) -> ProcessStatus: + """Wait for this process to reach a terminal state.""" + started_at = monotonic() + while True: + status = self.status + if status.done: + return status + if timeout is not None and monotonic() - started_at >= timeout: + raise TimeoutError( + f"Process {self.external_id or self.id!r} did not complete within " + f"{timeout} seconds" + ) + await asyncio.sleep(0.2) + + +class BaseRuntimeBundle(LoggingConfigurable): + """Compose compatible runtime, definition, and process implementations.""" + + runtime_class = Type(klass=BaseRuntime, default_value=BaseRuntime, config=True) + process_class = Type(klass=BaseProcess, default_value=BaseProcess, config=True) + definition_class = Type( + klass=BaseDefinition, + default_value=BaseDefinition, + config=True, + ) + # Services + default_dashboard_class = Type( + klass="beakerhub.services.dashboard.base.BaseDashboardService", + default_value="beakerhub.services.dashboard.base.BaseDashboardService", + config=True + ) + default_spawner_class = Type( + klass="beakerhub.services.spawner.base.BeakerhubImageSpawner", + default_value="beakerhub.services.spawner.base.BeakerhubImageSpawner", + config=True + ) + default_task_runner_class = Type( + klass="beakerhub.services.task.base.BaseTaskRunnerService", + default_value="beakerhub.services.task.base.BaseTaskRunnerService", + config=True + ) + + @property + def runtime(self): + """Resolve and return the runtime from the parent if defined.""" + if self.parent is not None: + return getattr(self.parent, "runtime", None) + + def create_definition( + self, + *, + runtime: BaseRuntime | None = None, + **kwargs: Any, + ) -> BaseDefinition: + """Create a definition using this bundle's provider implementation.""" + runtime = runtime or self.runtime + if runtime is not None: + if not isinstance(runtime, self.runtime_class): + raise TypeError( + f"{self.__class__.__name__} requires " + f"{self.runtime_class.__name__}, not {type(runtime).__name__}" + ) + kwargs["runtime"] = runtime + return self.definition_class(**kwargs) + + def start_process( + self, + definition: BaseDefinition, + *, + runtime: BaseRuntime | None = None, + process_type: ProcessType = "task", + **kwargs: Any, + ) -> BaseProcess: + """Launch a provider process using the application-owned runtime.""" + runtime = runtime or self.runtime + if not isinstance(definition, self.definition_class): + raise TypeError( + f"{self.__class__.__name__} requires " + f"{self.definition_class.__name__}, not {type(definition).__name__}" + ) + if not isinstance(runtime, self.runtime_class): + raise TypeError( + f"{self.__class__.__name__} requires " + f"{self.runtime_class.__name__}, not {type(runtime).__name__}" + ) + return self.process_class.start( + definition, + process_type=process_type, + runtime=runtime, + parent=runtime, + **kwargs, + ) diff --git a/src/beakerhub/runtimes/kubernetes.py b/src/beakerhub/runtimes/kubernetes.py new file mode 100644 index 0000000..6072acd --- /dev/null +++ b/src/beakerhub/runtimes/kubernetes.py @@ -0,0 +1,349 @@ +"""Kubernetes runtime workloads. + +This module owns Kubernetes Job mechanics shared by finite container workloads. +KubeSpawner remains responsible for JupyterHub session Pod construction and +proxy readiness until it is explicitly adapted to this runtime layer. +""" + +from typing import Any, Mapping +from uuid import uuid4 + +from kubernetes import client as k8s_client +from kubernetes import config as k8s_config +from traitlets import Dict, Instance, List, Unicode +import traitlets + +from beakerhub.runtimes.base import ( + BaseDefinition, + BaseProcess, + BaseRuntime, + BaseRuntimeBundle, + ProcessOutput, + ProcessStatus, + ProcessType, +) + + +class KubernetesRuntime(BaseRuntime): + """Shared Kubernetes clients and namespace defaults.""" + + namespace = Unicode( + "beakerhub", + config=True, + help="Kubernetes namespace for runtime workloads.", + ) + node_selector: dict[str, str] = Dict( + Unicode(), + Unicode(), + default_value={}, + config=True, + help="Default node selector for runtime workloads.", + ) + tolerations: list[dict[str, Any]] = List( + Dict(), + default_value=[], + config=True, + help="Default tolerations for runtime workloads.", + ) + service_account: str = Unicode( + "", + config=True, + help="Default service account for runtime workload Pods.", + ) + base_labels: dict[str, str] = Dict( + Unicode(), + Unicode(), + default_value={}, + config=True, + help="Labels applied to every runtime workload.", + ) + + batch_api: k8s_client.BatchV1Api + core_api: k8s_client.CoreV1Api + + + def __init__(self, **kwargs): + try: + k8s_config.load_incluster_config() + except k8s_config.ConfigException: + k8s_config.load_kube_config() + self.batch_api = k8s_client.BatchV1Api() + self.core_api = k8s_client.CoreV1Api() + super().__init__(**kwargs) + + @staticmethod + def build_resource_requirements( + resources: Mapping[str, Any], + ) -> k8s_client.V1ResourceRequirements | None: + """Convert request/limit mappings to Kubernetes resource requirements.""" + if not resources: + return None + return k8s_client.V1ResourceRequirements( + requests=resources.get("requests"), + limits=resources.get("limits"), + ) + + @staticmethod + def container_failure_message(status: Any, label: str) -> str | None: + """Return a useful message for a failed or waiting container.""" + if status.state and status.state.waiting: + reason = status.state.waiting.reason or "Unknown" + message = status.state.waiting.message or "" + return f"{label} '{status.name}' waiting: {reason}. {message}".strip() + if ( + status.state + and status.state.terminated + and status.state.terminated.exit_code != 0 + ): + reason = status.state.terminated.reason or "Error" + message = status.state.terminated.message or "" + return ( + f"{label} '{status.name}' failed ({reason}, exit code " + f"{status.state.terminated.exit_code}). {message}" + ).strip() + return None + + +class KubernetesDefinition(BaseDefinition): + """A Kubernetes Job workload definition.""" + + runtime = Instance(KubernetesRuntime, allow_none=True) + namespace = Unicode(config=True) + resources = Dict(default_value={}, config=True) + node_selector = Dict(Unicode(), Unicode(), config=True) + tolerations = List(Dict(), config=True) + service_account = Unicode(config=True) + backoff_limit = traitlets.Int(0, config=True) + active_deadline_seconds = traitlets.Int(300, allow_none=True, config=True) + ttl_seconds_after_finished = traitlets.Int(600, allow_none=True, config=True) + name_prefix = Unicode("beaker-process", config=True) + + @traitlets.default("namespace") + def _default_namespace(self) -> str: + return self.runtime.namespace + + @traitlets.default("node_selector") + def _default_node_selector(self) -> dict[str, str]: + return dict(self.runtime.node_selector) + + @traitlets.default("tolerations") + def _default_tolerations(self) -> list[dict[str, Any]]: + return list(self.runtime.tolerations) + + @traitlets.default("service_account") + def _default_service_account(self) -> str: + return self.runtime.service_account + + +class KubernetesProcess(BaseProcess): + """A Kubernetes Job-backed process.""" + + definition: KubernetesDefinition + runtime = Instance(KubernetesRuntime, allow_none=False) + + def __init__(self, definition: KubernetesDefinition, **kwargs: Any) -> None: + if not isinstance(definition, KubernetesDefinition): + raise TypeError("KubernetesProcess requires a KubernetesDefinition") + super().__init__(definition, **kwargs) + + @classmethod + def start( + cls, + definition: KubernetesDefinition, + *, + process_type: ProcessType = "task", + **kwargs: Any, + ) -> "KubernetesProcess": + """Create a Kubernetes Job and return its process handle.""" + if process_type != "task": + raise NotImplementedError( + "KubernetesProcess currently launches run-to-completion Jobs only" + ) + self = cls(definition, process_type=process_type, **kwargs) + self.external_id = self._job_name() + self.runtime.batch_api.create_namespaced_job( + namespace=self.definition.namespace, + body=self._job(), + ) + self.log.info("Created Kubernetes Job %s", self.external_id) + return self + + def _job_name(self) -> str: + prefix = self.definition.name_prefix.rstrip("-") or "beaker-process" + return f"{prefix}-{uuid4().hex[:8]}" + + def _job(self) -> k8s_client.V1Job: + definition = self.definition + labels = { + **self.runtime.base_labels, + **definition.labels, + "app.kubernetes.io/name": "beakerhub", + "app.kubernetes.io/component": "task", + "beakerhub/process-type": self.process_type, + } + container = k8s_client.V1Container( + name="task", + image=definition.image, + command=list(definition.entrypoint) or None, + args=list(definition.command) or None, + working_dir=definition.working_directory, + env=[ + k8s_client.V1EnvVar(name=name, value=value) + for name, value in definition.environment.items() + ] + or None, + resources=self.runtime.build_resource_requirements(definition.resources), + ) + pod_spec = k8s_client.V1PodSpec( + containers=[container], + restart_policy="Never", + node_selector=dict(definition.node_selector) or None, + tolerations=( + [k8s_client.V1Toleration(**item) for item in definition.tolerations] + if definition.tolerations + else None + ), + service_account_name=definition.service_account or None, + ) + return k8s_client.V1Job( + api_version="batch/v1", + kind="Job", + metadata=k8s_client.V1ObjectMeta( + name=self.external_id, + namespace=definition.namespace, + labels=labels, + ), + spec=k8s_client.V1JobSpec( + template=k8s_client.V1PodTemplateSpec( + metadata=k8s_client.V1ObjectMeta(labels=labels), + spec=pod_spec, + ), + backoff_limit=definition.backoff_limit, + active_deadline_seconds=definition.active_deadline_seconds, + ttl_seconds_after_finished=definition.ttl_seconds_after_finished, + ), + ) + + def describe(self) -> ProcessStatus: + """Return normalized lifecycle status for this Kubernetes Job.""" + if not self.external_id: + return ProcessStatus("pending", "Kubernetes Job has not been submitted") + try: + job = self.runtime.batch_api.read_namespaced_job( + name=self.external_id, + namespace=self.definition.namespace, + ) + except k8s_client.ApiException as error: + if error.status == 404: + return ProcessStatus("failed", f"Job {self.external_id} not found") + raise + + status = job.status + if status.succeeded and status.succeeded > 0: + return ProcessStatus("completed", "Job completed successfully") + if status.failed and status.failed > 0: + return ProcessStatus( + "failed", + self._failure_message(self.runtime.core_api), + ) + if status.active and status.active > 0: + return ProcessStatus("running", "Job is running") + return ProcessStatus("pending", "Job is pending") + + def _failure_message(self, core_api: k8s_client.CoreV1Api) -> str: + try: + pods = core_api.list_namespaced_pod( + namespace=self.definition.namespace, + label_selector=f"job-name={self.external_id}", + ) + except k8s_client.ApiException: + return "Job failed (could not retrieve pod details)" + if not pods.items: + return "Job failed (no pods found)" + pod = pods.items[0] + for status in pod.status.init_container_statuses or []: + message = self.runtime.container_failure_message(status, "Init container") + if message: + return message + for status in pod.status.container_statuses or []: + message = self.runtime.container_failure_message(status, "Container") + if message: + return message + phase = pod.status.phase or "Unknown" + reason = pod.status.reason or "" + return f"Job failed (pod phase: {phase}). {reason}".strip() + + def collect_output(self) -> ProcessOutput | None: + """Return the combined Kubernetes container log for this Job.""" + if not self.external_id: + return None + pods = self.runtime.core_api.list_namespaced_pod( + namespace=self.definition.namespace, + label_selector=f"job-name={self.external_id}", + ) + if not pods.items: + return None + response = self.runtime.core_api.read_namespaced_pod_log( + name=pods.items[0].metadata.name, + namespace=self.definition.namespace, + container="task", + _preload_content=False, + ) + logs = response.data.decode("utf-8") + return ProcessOutput(stdout=logs or "", stderr="") + + def stop(self) -> None: + """Delete this Kubernetes Job and its associated Pods.""" + if not self.external_id: + return + try: + self.runtime.batch_api.delete_namespaced_job( + name=self.external_id, + namespace=self.definition.namespace, + body=k8s_client.V1DeleteOptions(propagation_policy="Background"), + ) + except k8s_client.ApiException as error: + if error.status != 404: + raise + + +class KubernetesRuntimeBundle(BaseRuntimeBundle): + """Composition bundle for Kubernetes Job workloads.""" + + runtime_class = traitlets.Type( + klass=KubernetesRuntime, + default_value=KubernetesRuntime, + config=True, + ) + process_class = traitlets.Type( + klass=KubernetesProcess, + default_value=KubernetesProcess, + config=True, + ) + definition_class = traitlets.Type( + klass=KubernetesDefinition, + default_value=KubernetesDefinition, + config=True, + ) + + default_dashboard_class = traitlets.Type( + klass="beakerhub.services.dashboard.kubernetes_dashboard.KubernetesDashboardService", + default_value="beakerhub.services.dashboard.kubernetes_dashboard.KubernetesDashboardService", + config=True + ) + default_spawner_class = traitlets.Type( + klass="beakerhub.services.spawner.kubernetes_spawner.BeakerKubeSpawner", + default_value="beakerhub.services.spawner.kubernetes_spawner.BeakerKubeSpawner", + config=True + ) + default_task_runner_class = traitlets.Type( + klass="beakerhub.services.task.kubernetes_task_runner.KubernetesTaskRunnerService", + default_value="beakerhub.services.task.kubernetes_task_runner.KubernetesTaskRunnerService", + config=True + ) + + + +# KubeSpawner owns session Pod creation, state, and proxy readiness today. A +# future session adapter can share KubernetesRuntime client/configuration logic +# without making Kubernetes Jobs imitate JupyterHub server semantics. diff --git a/src/beakerhub/services/dashboard/__init__.py b/src/beakerhub/services/dashboard/__init__.py new file mode 100644 index 0000000..bcffc8c --- /dev/null +++ b/src/beakerhub/services/dashboard/__init__.py @@ -0,0 +1,5 @@ +"""Dashboard data services.""" + +from .base import BaseDashboardService, DashboardServiceError + +__all__ = ["BaseDashboardService", "DashboardServiceError"] diff --git a/src/beakerhub/services/dashboard/aws_ecs_dashboard.py b/src/beakerhub/services/dashboard/aws_ecs_dashboard.py new file mode 100644 index 0000000..195df44 --- /dev/null +++ b/src/beakerhub/services/dashboard/aws_ecs_dashboard.py @@ -0,0 +1,322 @@ +"""AWS ECS implementation of the dashboard service.""" + +from datetime import datetime, timezone +from typing import Any + +from botocore.exceptions import ClientError + +from beakerhub.runtimes.aws_ecs import AwsEcsRuntime +from beakerhub.services.dashboard.base import ( + BaseDashboardService, + DashboardServiceError, +) + + +class AwsEcsDashboardService(BaseDashboardService): + """Provide AWS ECS runtime data and session logs to the dashboard.""" + + def _runtime(self) -> AwsEcsRuntime: + parent_runtime = getattr(self.parent, "runtime", None) + if isinstance(parent_runtime, AwsEcsRuntime): + return parent_runtime + return AwsEcsRuntime(parent=self.parent) + + def get_dashboard(self) -> dict[str, Any]: + runtime = self._runtime() + cluster_name = runtime.cluster_name + if not cluster_name: + return { + "available": False, + "error": "No ECS cluster is configured", + } + + try: + cluster = runtime.ecs_client.describe_clusters( + clusters=[cluster_name], + ).get("clusters", [])[0] + except (ClientError, IndexError) as error: + return {"available": False, "error": str(error)} + + try: + services = self._describe_services(runtime, cluster_name) + tasks = self._describe_tasks(runtime, cluster_name) + stopped_tasks = self._describe_stopped_tasks(runtime, cluster_name) + instances = self._describe_container_instances(runtime, cluster_name) + except ClientError as error: + return {"available": False, "error": str(error)} + + is_fargate = self._is_fargate(runtime, tasks, instances) + running_tasks = sum(1 for task in tasks if task.get("lastStatus") == "RUNNING") + pending_tasks = sum(1 for task in tasks if task.get("lastStatus") == "PENDING") + failed_tasks = self._failed_tasks(stopped_tasks) + task_detail = f"{running_tasks} running, {pending_tasks} pending" + if failed_tasks: + task_detail += f", {len(failed_tasks)} failed" + return { + "available": True, + "runtime": { + "provider": "Amazon ECS", + "scope": self._short_name(cluster.get("clusterArn", cluster_name)), + }, + "summary": [ + { + "label": "Services", + "value": len(services), + "detail": f"{sum(service.get('runningCount', 0) for service in services)} running", + }, + { + "label": "Tasks", + "value": len(tasks), + "detail": task_detail, + "severity": "danger" if failed_tasks else None, + }, + { + "label": "Compute instances", + "value": len(instances), + "detail": ( + "Compute is managed by AWS Fargate. This cluster has no " + "registered container instances." + if is_fargate + else f"{sum(1 for item in instances if item.get('status') == 'ACTIVE')} active" + ), + }, + ], + "workloads": self._workloads(services, tasks), + "resources_empty_message": ( + "AWS Fargate manages the underlying compute resources. Per-instance " + "capacity is not available." + if is_fargate and not instances + else None + ), + "resources": [ + { + "name": self._short_name(instance.get("containerInstanceArn", "Unknown")), + "status": instance.get("status", "UNKNOWN").lower(), + "details": [ + {"label": "Agent", "value": instance.get("agentConnected", False) and "Connected" or "Disconnected"}, + {"label": "Running tasks", "value": instance.get("runningTasksCount", 0)}, + {"label": "Pending tasks", "value": instance.get("pendingTasksCount", 0)}, + ], + } + for instance in instances + ], + "alerts": [ + { + "reason": "Stopped task", + "message": task.get("stoppedReason", "Task stopped"), + "object": self._short_name(task.get("taskArn", "Unknown")), + "timestamp": task.get("stoppedAt", datetime.now(timezone.utc)).isoformat(), + } + for task in failed_tasks + ][:20], + } + + def get_session_logs( + self, + session_id: str, + container: str, + tail_lines: int, + ) -> dict[str, Any]: + runtime = self._runtime() + if not runtime.cluster_name: + raise DashboardServiceError(503, "No ECS cluster is configured") + task = runtime.describe_task(runtime.cluster_name, session_id) + if task is None: + raise DashboardServiceError(404, "Session task was not found") + + task_id = session_id.rsplit("/", maxsplit=1)[-1] + task_container = next( + (item for item in task.get("containers", []) if item.get("name") == container), + None, + ) + if task_container is None and container == "notebook": + task_container = next(iter(task.get("containers", [])), None) + if task_container is None: + raise DashboardServiceError(404, f"Container '{container}' was not found") + container_name = task_container.get("name", container) + if not runtime.log_group: + raise DashboardServiceError(503, "CloudWatch log group is not configured") + + stream_name = task_container.get( + "logStreamName", + f"{runtime.log_stream_prefix}/{container_name}/{task_id}", + ) + try: + messages, truncated = runtime.get_log_tail( + runtime.log_group, + stream_name, + tail_lines, + ) + except RuntimeError as error: + raise DashboardServiceError(502, str(error)) from error + return { + "runtime_name": self._short_name(task.get("taskArn", session_id)), + "pod_name": self._short_name(task.get("taskArn", session_id)), + "container": container_name, + "logs": "\n".join(messages), + "tail_lines": tail_lines, + "truncated": truncated, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + + @staticmethod + def _failed_tasks(tasks: list[dict[str, Any]]) -> list[dict[str, Any]]: + failure_stop_codes = {"EssentialContainerExited", "TaskFailedToStart"} + return [ + task + for task in tasks + if task.get("stopCode") in failure_stop_codes + or ( + task.get("stopCode") != "UserInitiated" + and any( + container.get("exitCode", 0) != 0 + for container in task.get("containers", []) + ) + ) + ] + + @staticmethod + def _is_fargate( + runtime: AwsEcsRuntime, + tasks: list[dict[str, Any]], + instances: list[dict[str, Any]], + ) -> bool: + """Determine whether the displayed capacity is Fargate-managed.""" + if instances: + return False + providers = { + task.get("capacityProviderName") + for task in tasks + if task.get("capacityProviderName") + } + if providers: + return providers <= {"FARGATE", "FARGATE_SPOT"} + strategy = runtime.launch_options.get("capacityProviderStrategy", []) + if strategy: + return { + item.get("capacityProvider") for item in strategy + } <= {"FARGATE", "FARGATE_SPOT"} + return runtime.launch_type == "FARGATE" + + def _workloads( + self, + services: list[dict[str, Any]], + tasks: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Group service-owned tasks under their ECS service.""" + managed_task_arns: set[str] = set() + workloads = [] + for service in services: + service_name = service.get("serviceName", "Unknown service") + owned_tasks = [ + task + for task in tasks + if task.get("group") == f"service:{service_name}" + ] + children = [self._task_workload(task) for task in owned_tasks] + managed_task_arns.update(task.get("taskArn", "") for task in owned_tasks) + workloads.append( + { + "name": service_name, + "kind": "Service", + "status": service.get("status", "UNKNOWN").lower(), + "detail": ( + f"{service.get('runningCount', 0)} running / " + f"{service.get('desiredCount', 0)} desired" + ), + "children": children, + } + ) + workloads.extend( + self._task_workload(task) + for task in tasks + if task.get("taskArn", "") not in managed_task_arns + ) + return workloads + + def _task_workload(self, task: dict[str, Any]) -> dict[str, str]: + return { + "name": self._short_name(task.get("taskArn", "Unknown task")), + "kind": "Task", + "status": task.get("lastStatus", "UNKNOWN").lower(), + "detail": task.get("taskDefinitionArn", "").rsplit("/", maxsplit=1)[-1], + } + + @staticmethod + def _short_name(value: str) -> str: + return value.rsplit("/", maxsplit=1)[-1] + + @staticmethod + def _list_all(client_method, result_key: str, **kwargs: Any) -> list[str]: + values: list[str] = [] + next_token: str | None = None + while True: + request = dict(kwargs) + if next_token: + request["nextToken"] = next_token + response = client_method(**request) + values.extend(response.get(result_key, [])) + next_token = response.get("nextToken") + if not next_token: + return values + + def _describe_services(self, runtime: AwsEcsRuntime, cluster: str) -> list[dict[str, Any]]: + arns = self._list_all( + runtime.ecs_client.list_services, + "serviceArns", + cluster=cluster, + ) + return [ + service + for start in range(0, len(arns), 10) + for service in runtime.ecs_client.describe_services( + cluster=cluster, services=arns[start : start + 10] + ).get("services", []) + ] + + def _describe_tasks(self, runtime: AwsEcsRuntime, cluster: str) -> list[dict[str, Any]]: + arns = self._list_all( + runtime.ecs_client.list_tasks, + "taskArns", + cluster=cluster, + ) + return [ + task + for start in range(0, len(arns), 100) + for task in runtime.ecs_client.describe_tasks( + cluster=cluster, tasks=arns[start : start + 100] + ).get("tasks", []) + ] + + def _describe_stopped_tasks( + self, runtime: AwsEcsRuntime, cluster: str + ) -> list[dict[str, Any]]: + arns = self._list_all( + runtime.ecs_client.list_tasks, + "taskArns", + cluster=cluster, + desiredStatus="STOPPED", + ) + return [ + task + for start in range(0, len(arns), 100) + for task in runtime.ecs_client.describe_tasks( + cluster=cluster, tasks=arns[start : start + 100] + ).get("tasks", []) + ] + + def _describe_container_instances( + self, runtime: AwsEcsRuntime, cluster: str + ) -> list[dict[str, Any]]: + arns = self._list_all( + runtime.ecs_client.list_container_instances, + "containerInstanceArns", + cluster=cluster, + ) + return [ + instance + for start in range(0, len(arns), 100) + for instance in runtime.ecs_client.describe_container_instances( + cluster=cluster, containerInstances=arns[start : start + 100] + ).get("containerInstances", []) + ] diff --git a/src/beakerhub/services/dashboard/base.py b/src/beakerhub/services/dashboard/base.py new file mode 100644 index 0000000..dcdd818 --- /dev/null +++ b/src/beakerhub/services/dashboard/base.py @@ -0,0 +1,36 @@ +"""Base contract for dashboard data services.""" + +from typing import Any + +from traitlets.config import LoggingConfigurable + + +class DashboardServiceError(Exception): + """An error that a dashboard handler can expose as an HTTP response.""" + + def __init__(self, status_code: int, reason: str): + super().__init__(reason) + self.status_code = status_code + self.reason = reason + + +class BaseDashboardService(LoggingConfigurable): + """Provide runtime inventory and session logs to dashboard handlers.""" + + def get_dashboard(self) -> dict[str, Any]: + """Return normalized runtime inventory for the admin dashboard. + + Implementations return ``available`` and, when available, ``runtime``, + ``summary``, ``workloads``, ``resources``, and ``alerts``. The values + describe container-runtime concepts rather than provider API objects. + """ + raise NotImplementedError + + def get_session_logs( + self, + session_id: str, + container: str, + tail_lines: int, + ) -> dict[str, Any]: + """Return logs for a session runtime.""" + raise NotImplementedError diff --git a/src/beakerhub/services/dashboard/handlers.py b/src/beakerhub/services/dashboard/handlers.py new file mode 100644 index 0000000..800436a --- /dev/null +++ b/src/beakerhub/services/dashboard/handlers.py @@ -0,0 +1,252 @@ +"""API handlers for BeakerHub dashboard services.""" + +import json +import logging + +from jupyterhub.apihandlers import APIHandler +from jupyterhub.orm import Spawner, User +from jupyterhub.scopes import needs_scope +from sqlalchemy import func +from tornado import web + +from beakerhub import __version__ +from beakerhub.orm import ( + BeakerSession, + Context, + Integration, + Language, + NodeImages, + NodeImageTask, + NodeSecret, + Workflow, +) +from beakerhub.services.dashboard.base import DashboardServiceError + + +log = logging.getLogger(__name__) + + +class AdminDashboardSummaryHandler(APIHandler): + """Return application-level counts for the admin dashboard.""" + + def compute_etag(self) -> None: + return None + + @needs_scope("admin:users") + async def get(self): + db = self.db + images_total = db.query(func.count(NodeImages.id)).scalar() or 0 + images_enabled = ( + db.query(func.count(NodeImages.id)) + .filter(NodeImages.enabled.is_(True)) + .scalar() + or 0 + ) + contexts_total = db.query(func.count(Context.id)).scalar() or 0 + contexts_enabled = ( + db.query(func.count(Context.id)) + .filter(Context.enabled.is_(True)) + .scalar() + or 0 + ) + workflows_total = db.query(func.count(Workflow.id)).scalar() or 0 + workflows_enabled = ( + db.query(func.count(Workflow.id)) + .filter(Workflow.enabled.is_(True)) + .scalar() + or 0 + ) + integrations_total = db.query(func.count(Integration.id)).scalar() or 0 + integrations_enabled = ( + db.query(func.count(Integration.id)) + .filter(Integration.enabled.is_(True)) + .scalar() + or 0 + ) + languages_total = db.query(func.count(Language.slug)).scalar() or 0 + secrets_total = db.query(func.count(NodeSecret.id)).scalar() or 0 + secrets_global = ( + db.query(func.count(NodeSecret.id)) + .filter(NodeSecret.node_image_id.is_(None)) + .scalar() + or 0 + ) + users_total = db.query(func.count(User.id)).scalar() or 0 + users_admin = ( + db.query(func.count(User.id)).filter(User.admin.is_(True)).scalar() or 0 + ) + active_servers = ( + db.query(func.count(Spawner.id)) + .filter(Spawner.server_id.isnot(None)) + .scalar() + or 0 + ) + beaker_sessions_total = db.query(func.count(BeakerSession.id)).scalar() or 0 + recent_tasks = ( + db.query(NodeImageTask) + .join(NodeImages, NodeImageTask.node_image_id == NodeImages.id) + .order_by(NodeImageTask.updated_at.desc()) + .limit(5) + .all() + ) + recent_imports = [ + { + "task_id": task.id, + "node_image_slug": ( + task.node_image.slug if task.node_image else None + ), + "task_type": task.task_type, + "status": task.status, + "job_name": task.job_name, + "created_at": ( + task.created_at.isoformat() if task.created_at else None + ), + "updated_at": ( + task.updated_at.isoformat() if task.updated_at else None + ), + "error": task.error, + "result": task.result, + } + for task in recent_tasks + ] + result = { + "app_version": __version__, + "counts": { + "images": {"total": images_total, "enabled": images_enabled}, + "contexts": {"total": contexts_total, "enabled": contexts_enabled}, + "workflows": {"total": workflows_total, "enabled": workflows_enabled}, + "integrations": { + "total": integrations_total, + "enabled": integrations_enabled, + }, + "languages": {"total": languages_total}, + "secrets": { + "total": secrets_total, + "global": secrets_global, + "per_node": secrets_total - secrets_global, + }, + "users": {"total": users_total, "admin": users_admin}, + "sessions": { + "active_servers": active_servers, + "beaker_sessions_total": beaker_sessions_total, + }, + }, + "recent_imports": recent_imports, + } + self.set_header("Content-Type", "application/json") + self.write(json.dumps(result)) + + +class AdminDashboardClusterHandler(APIHandler): + """Return data from the configured dashboard service.""" + + def compute_etag(self) -> None: + return None + + @needs_scope("admin:users") + async def get(self): + try: + result = self.settings["app"].dashboard_service.get_dashboard() + self._add_task_sessions(result) + except Exception as error: + log.warning("Failed to fetch dashboard data: %s", error) + result = {"available": False, "error": str(error)} + self.set_header("Content-Type", "application/json") + self.write(json.dumps(result)) + + + def _add_task_sessions(self, dashboard: dict) -> None: + """Attach JupyterHub session ownership to task workload rows. + + ECS session spawners persist their task ARN in state. Matching that ARN + avoids one ECS tag request per task and also identifies the session owner. + """ + task_sessions: dict[str, list[dict[str, str]]] = {} + for spawner in ( + self.db.query(Spawner) + .join(User) + .filter(Spawner.server_id.isnot(None)) + .all() + ): + state = spawner.state if isinstance(spawner.state, dict) else {} + task_arn = state.get("task_arn") + if not task_arn: + continue + task_id = task_arn.rsplit("/", maxsplit=1)[-1] + task_sessions.setdefault(task_id, []).append( + {"user": spawner.user.name, "name": spawner.name} + ) + + def add_sessions(workload: dict) -> None: + if workload.get("kind") == "Task": + sessions = task_sessions.get(workload.get("name", "")) + if sessions: + workload["sessions"] = sessions + for child in workload.get("children", []): + add_sessions(child) + + for workload in dashboard.get("workloads", []): + add_sessions(workload) + + +class AdminSessionLogsHandler(APIHandler): + """Return logs from the configured dashboard service.""" + + def compute_etag(self) -> None: + return None + + @needs_scope("admin:users") + async def get(self, owner: str, session_id: str): + container = self.get_argument("container", "notebook") + try: + tail_lines = int(self.get_argument("tail_lines", "5000")) + except ValueError: + tail_lines = 5000 + tail_lines = max(1, min(tail_lines, 100000)) + + runtime_id = session_id + spawner = ( + self.db.query(Spawner) + .join(User) + .filter(User.name == owner, Spawner.name == session_id) + .first() + ) + if spawner and isinstance(spawner.state, dict): + runtime_id = spawner.state.get("task_arn", session_id) + + try: + result = self.settings["app"].dashboard_service.get_session_logs( + runtime_id, + container, + tail_lines, + ) + except DashboardServiceError as error: + raise web.HTTPError(error.status_code, reason=error.reason) from error + except Exception as error: + log.warning( + "Failed to fetch session logs for %s/%s: %s", + owner, + session_id, + error, + ) + raise web.HTTPError(500, reason=str(error)) from error + + self.set_header("Content-Type", "application/json") + self.write(json.dumps(result)) + + +handlers = [ + (r"/api/beakerhub/admin/dashboard/summary", AdminDashboardSummaryHandler), + (r"/api/beakerhub/admin/dashboard/cluster", AdminDashboardClusterHandler), + ( + r"/api/beakerhub/admin/dashboard/session-logs/([^/]+)/([^/]+)", + AdminSessionLogsHandler, + ), + # Compatibility path for existing admin clients. + ( + r"/api/beakerhub/admin/dashboard/pod-logs/([^/]+)/([^/]+)", + AdminSessionLogsHandler, + ), +] + +api_handlers = handlers diff --git a/src/beakerhub/dashboard_handlers.py b/src/beakerhub/services/dashboard/kubernetes_dashboard.py similarity index 64% rename from src/beakerhub/dashboard_handlers.py rename to src/beakerhub/services/dashboard/kubernetes_dashboard.py index f2e280a..d08a057 100644 --- a/src/beakerhub/dashboard_handlers.py +++ b/src/beakerhub/services/dashboard/kubernetes_dashboard.py @@ -1,170 +1,43 @@ -"""Admin dashboard API handlers for BeakerHub. - -Provides aggregated summary statistics and optional Kubernetes cluster -information for the admin dashboard page. -""" -import json +"""Kubernetes implementation of the dashboard service.""" import logging from datetime import datetime, timezone, timedelta from typing import Any -from tornado import web -from jupyterhub.apihandlers import APIHandler -from jupyterhub.scopes import needs_scope -from jupyterhub.orm import User, Spawner -from sqlalchemy import func - -from beakerhub import __version__ -from beakerhub.orm import ( - BeakerSession, - Context, - Integration, - Language, - NodeImages, - NodeImageTask, - NodeSecret, - Workflow, +from traitlets import Unicode + +from beakerhub.runtimes.kubernetes import KubernetesRuntime +from beakerhub.services.dashboard.base import ( + BaseDashboardService, + DashboardServiceError, ) log = logging.getLogger(__name__) -class AdminDashboardSummaryHandler(APIHandler): - """Aggregated counts and summary statistics for the admin dashboard.""" - - def compute_etag(self) -> None: - return None - - @needs_scope('admin:users') - async def get(self): - db = self.db - - # Entity counts - images_total = db.query(func.count(NodeImages.id)).scalar() or 0 - images_enabled = db.query(func.count(NodeImages.id)).filter( - NodeImages.enabled == True - ).scalar() or 0 - - contexts_total = db.query(func.count(Context.id)).scalar() or 0 - contexts_enabled = db.query(func.count(Context.id)).filter( - Context.enabled == True - ).scalar() or 0 - - workflows_total = db.query(func.count(Workflow.id)).scalar() or 0 - workflows_enabled = db.query(func.count(Workflow.id)).filter( - Workflow.enabled == True - ).scalar() or 0 - - integrations_total = db.query(func.count(Integration.id)).scalar() or 0 - integrations_enabled = db.query(func.count(Integration.id)).filter( - Integration.enabled == True - ).scalar() or 0 - - languages_total = db.query(func.count(Language.slug)).scalar() or 0 - - secrets_total = db.query(func.count(NodeSecret.id)).scalar() or 0 - secrets_global = db.query(func.count(NodeSecret.id)).filter( - NodeSecret.node_image_id.is_(None) - ).scalar() or 0 - - # User counts - users_total = db.query(func.count(User.id)).scalar() or 0 - users_admin = db.query(func.count(User.id)).filter( - User.admin == True - ).scalar() or 0 - - # Active sessions: JupyterHub spawners with a server_id (meaning running) - active_servers = db.query(func.count(Spawner.id)).filter( - Spawner.server_id.isnot(None) - ).scalar() or 0 - - # BeakerSession counts (historical) - beaker_sessions_total = db.query(func.count(BeakerSession.id)).scalar() or 0 - - # Recent import tasks (last 5) - recent_tasks = db.query(NodeImageTask).join( - NodeImages, NodeImageTask.node_image_id == NodeImages.id - ).order_by( - NodeImageTask.updated_at.desc() - ).limit(5).all() - - recent_imports = [] - for task in recent_tasks: - recent_imports.append({ - "task_id": task.id, - "node_image_slug": task.node_image.slug if task.node_image else None, - "task_type": task.task_type, - "status": task.status, - "job_name": task.job_name, - "created_at": task.created_at.isoformat() if task.created_at else None, - "updated_at": task.updated_at.isoformat() if task.updated_at else None, - "error": task.error, - "result": task.result, - }) - - result = { - "app_version": __version__, - "counts": { - "images": {"total": images_total, "enabled": images_enabled}, - "contexts": {"total": contexts_total, "enabled": contexts_enabled}, - "workflows": {"total": workflows_total, "enabled": workflows_enabled}, - "integrations": {"total": integrations_total, "enabled": integrations_enabled}, - "languages": {"total": languages_total}, - "secrets": {"total": secrets_total, "global": secrets_global, "per_node": secrets_total - secrets_global}, - "users": {"total": users_total, "admin": users_admin}, - "sessions": { - "active_servers": active_servers, - "beaker_sessions_total": beaker_sessions_total, - }, - }, - "recent_imports": recent_imports, - } - - self.set_header("Content-Type", "application/json") - self.write(json.dumps(result)) - - -class AdminDashboardClusterHandler(APIHandler): - """Kubernetes cluster information for the admin dashboard. +class KubernetesDashboardService(BaseDashboardService): + """Provide Kubernetes cluster data and Pod logs.""" - Queries the K8s API for pod, PVC, event, and job information. - Returns available: false if K8s access is not available. - """ + namespace = Unicode( + "beakerhub", + config=True, + help="Kubernetes namespace inspected by the dashboard service.", + ) - def compute_etag(self) -> None: - return None + def _runtime(self) -> KubernetesRuntime: + parent_runtime = getattr(self.parent, "runtime", None) + if isinstance(parent_runtime, KubernetesRuntime): + return parent_runtime + return KubernetesRuntime(namespace=self.namespace) - @needs_scope('admin:users') - async def get(self): - try: - result = self._get_cluster_info() - except Exception as e: - log.warning(f"Failed to fetch K8s cluster info: {e}") - result = { - "available": False, - "error": str(e), - } - - self.set_header("Content-Type", "application/json") - self.write(json.dumps(result)) - - def _get_cluster_info(self) -> dict[str, Any]: - from kubernetes import client as k8s_client + def get_dashboard(self) -> dict[str, Any]: from kubernetes import config as k8s_config try: - k8s_config.load_incluster_config() + batch_api, core_api = self._runtime().get_clients() except k8s_config.ConfigException: - try: - k8s_config.load_kube_config() - except k8s_config.ConfigException: - return {"available": False, "error": "No K8s configuration found"} - - core_api = k8s_client.CoreV1Api() - batch_api = k8s_client.BatchV1Api() + return {"available": False, "error": "No K8s configuration found"} - app = self.settings.get("app") - namespace = getattr(app, "task_namespace", "beakerhub") + namespace = self.namespace # Pods pods_result = self._get_pod_info(core_api, namespace) @@ -181,20 +54,85 @@ def _get_cluster_info(self) -> dict[str, Any]: # Cluster nodes (cluster-scoped — requires ClusterRole) nodes_result = self._get_node_info(core_api) - # Helm releases (stored as secrets with owner=helm label) - helm_result = self._get_helm_releases(core_api, namespace) - + pod_phases = pods_result.get("by_phase", {}) return { "available": True, - "namespace": namespace, - "pods": pods_result, - "pvcs": pvcs_result, - "jobs": jobs_result, - "events": events_result, - "nodes": nodes_result, - "helm_releases": helm_result, + "runtime": {"provider": "Kubernetes", "scope": namespace}, + "summary": [ + { + "label": "Workloads", + "value": pods_result.get("total", 0), + "detail": f"{pod_phases.get('Running', 0)} running, " + f"{pod_phases.get('Pending', 0)} pending", + }, + { + "label": "Tasks", + "value": jobs_result.get("total", 0), + "detail": f"{jobs_result.get('active', 0)} active, " + f"{jobs_result.get('failed', 0)} failed", + "severity": "danger" if jobs_result.get("failed", 0) else None, + }, + { + "label": "Storage volumes", + "value": len(pvcs_result), + "detail": f"{sum(1 for pvc in pvcs_result if pvc.get('phase') == 'Bound')} ready", + }, + ], + "workloads": [ + { + "name": name, + "kind": "Workload", + "status": self._component_status(info.get("phases", {})), + "detail": f"{info.get('count', 0)} instances", + } + for name, info in pods_result.get("by_component", {}).items() + ], + "resources": [ + { + "name": node.get("name", "Unknown"), + "status": "ready" if node.get("ready") else "unavailable", + "details": [ + {"label": "Instance", "value": node.get("instance_type") or "—"}, + {"label": "CPU", "value": f"{node.get('allocated', {}).get('cpu', '0')} / {node.get('allocatable', {}).get('cpu', '?')}"}, + {"label": "Memory", "value": f"{node.get('allocated', {}).get('memory', '0')} / {node.get('allocatable', {}).get('memory', '?')}"}, + ], + } + for node in nodes_result + if not node.get("error") + ] + [ + { + "name": pvc.get("name", "Unknown"), + "status": (pvc.get("phase") or "unknown").lower(), + "details": [ + {"label": "Storage", "value": pvc.get("capacity") or "Unknown"}, + {"label": "Class", "value": pvc.get("storage_class") or "Default"}, + ], + } + for pvc in pvcs_result + if not pvc.get("error") + ], + "alerts": [ + { + "reason": event.get("reason", "Warning"), + "message": event.get("message", ""), + "object": event.get("involved_object"), + "timestamp": event.get("last_timestamp"), + } + for event in events_result + if not event.get("error") + ], } + @staticmethod + def _component_status(phases: dict[str, int]) -> str: + if phases.get("Failed"): + return "failed" + if phases.get("Pending"): + return "pending" + if phases.get("Running"): + return "running" + return "unknown" + def _get_pod_info(self, core_api, namespace: str) -> dict[str, Any]: try: pods = core_api.list_namespaced_pod(namespace=namespace) @@ -509,52 +447,21 @@ def _get_helm_releases(self, core_api, namespace: str) -> list[dict[str, Any]]: return sorted(result, key=lambda r: r.get("name", "")) -class AdminPodLogsHandler(APIHandler): - """Fetch container logs for a session pod. - - Locates the pod via the ``hub.jupyter.org/servername`` label and reads - logs from the specified container (default: ``notebook``). - """ - - def compute_etag(self) -> None: - return None - - @needs_scope('admin:users') - async def get(self, owner: str, session_id: str): - container = self.get_argument("container", "notebook") - try: - tail_lines = int(self.get_argument("tail_lines", "5000")) - except ValueError: - tail_lines = 5000 - tail_lines = max(1, min(tail_lines, 100000)) - - try: - result = self._get_pod_logs(session_id, container, tail_lines) - except Exception as e: - log.warning(f"Failed to fetch pod logs for {owner}/{session_id}: {e}") - raise web.HTTPError(500, reason=str(e)) - - self.set_header("Content-Type", "application/json") - self.write(json.dumps(result)) - - def _get_pod_logs( + def get_session_logs( self, server_name: str, container: str, tail_lines: int ) -> dict[str, Any]: from kubernetes import client as k8s_client from kubernetes import config as k8s_config try: - k8s_config.load_incluster_config() + _, core_api = self._runtime().get_clients() except k8s_config.ConfigException: - try: - k8s_config.load_kube_config() - except k8s_config.ConfigException: - raise web.HTTPError(503, reason="No Kubernetes configuration found") - - core_api = k8s_client.CoreV1Api() + raise DashboardServiceError( + 503, + "No Kubernetes configuration found", + ) - app = self.settings.get("app") - namespace = getattr(app, "task_namespace", "beakerhub") + namespace = self.namespace # Find the pod by JupyterHub server-name label label_selector = f"hub.jupyter.org/servername={server_name}" @@ -562,8 +469,9 @@ def _get_pod_logs( namespace=namespace, label_selector=label_selector ) if not pods.items: - raise web.HTTPError( - 404, reason=f"No pod found with servername={server_name}" + raise DashboardServiceError( + 404, + f"No pod found with servername={server_name}", ) pod = pods.items[0] @@ -578,9 +486,9 @@ def _get_pod_logs( ) except k8s_client.exceptions.ApiException as e: if e.status == 400: - raise web.HTTPError( + raise DashboardServiceError( 400, - reason=f"Container '{container}' not found in pod '{pod_name}'", + f"Container '{container}' not found in pod '{pod_name}'", ) raise @@ -588,6 +496,7 @@ def _get_pod_logs( truncated = len(log_lines) >= tail_lines return { + "runtime_name": pod_name, "pod_name": pod_name, "container": container, "logs": logs or "", @@ -595,10 +504,3 @@ def _get_pod_logs( "truncated": truncated, "timestamp": datetime.now(timezone.utc).isoformat(), } - - -dashboard_handlers = [ - (r"/api/beakerhub/admin/dashboard/summary", AdminDashboardSummaryHandler), - (r"/api/beakerhub/admin/dashboard/cluster", AdminDashboardClusterHandler), - (r"/api/beakerhub/admin/dashboard/pod-logs/([^/]+)/([^/]+)", AdminPodLogsHandler), -] diff --git a/src/beakerhub/services/periodic_tasks/__init__.py b/src/beakerhub/services/periodic_tasks/__init__.py new file mode 100644 index 0000000..3b29171 --- /dev/null +++ b/src/beakerhub/services/periodic_tasks/__init__.py @@ -0,0 +1 @@ +"""Periodic managed services.""" diff --git a/src/beakerhub/services/idle_culler/__init__.py b/src/beakerhub/services/periodic_tasks/idle_culler/__init__.py similarity index 89% rename from src/beakerhub/services/idle_culler/__init__.py rename to src/beakerhub/services/periodic_tasks/idle_culler/__init__.py index 6b7281d..f65d236 100644 --- a/src/beakerhub/services/idle_culler/__init__.py +++ b/src/beakerhub/services/periodic_tasks/idle_culler/__init__.py @@ -7,7 +7,7 @@ "command": [ sys.executable, "-m", - "beakerhub.services.idle_culler.idle_culler_service", + "beakerhub.services.periodic_tasks.idle_culler.service", ], } diff --git a/src/beakerhub/services/idle_culler/idle_culler_service.py b/src/beakerhub/services/periodic_tasks/idle_culler/service.py similarity index 100% rename from src/beakerhub/services/idle_culler/idle_culler_service.py rename to src/beakerhub/services/periodic_tasks/idle_culler/service.py diff --git a/src/beakerhub/services/spawner/__init__.py b/src/beakerhub/services/spawner/__init__.py new file mode 100644 index 0000000..80d6fe7 --- /dev/null +++ b/src/beakerhub/services/spawner/__init__.py @@ -0,0 +1,6 @@ +"""Spawner services for BeakerHub session runtimes.""" + +from .aws_ecs_spawner import BeakerAwsECSSpawner +from .base import BeakerSpawner, BeakerhubImageSpawner + +__all__ = ["BeakerAwsECSSpawner", "BeakerSpawner", "BeakerhubImageSpawner"] diff --git a/src/beakerhub/services/spawner/aws_ecs_spawner.py b/src/beakerhub/services/spawner/aws_ecs_spawner.py new file mode 100644 index 0000000..d47fe77 --- /dev/null +++ b/src/beakerhub/services/spawner/aws_ecs_spawner.py @@ -0,0 +1,613 @@ +"""AWS ECS-backed JupyterHub spawner.""" + +import asyncio +import hashlib +import os +import time +from functools import cache +from typing import Any, Literal + +import boto3 +import traitlets +from kubespawner.slugs import strip_and_hash +from traitlets import Bool, Dict, Integer, Instance, List, Unicode, default +from traitlets.config import Config + +from beakerhub import orm +from beakerhub.runtimes.aws_ecs import AwsEcsDefinition, AwsEcsProcess, AwsEcsRuntime +from beakerhub.runtimes.base import ProcessStatus +from beakerhub.services.spawner.base import BeakerhubImageSpawner + + +_CONFIG_VOLUME_NAME = "config" +_CONFIG_WRITER_NAME = "config-writer" +_CONFIG_WRITER_IMAGE = "busybox:latest" +_CONFIG_WRITER_MOUNT_PATH = "/opt/beaker/config" +_NOTEBOOK_CONFIG_PATH = "/root/.config" + + +class BeakerAwsECSSpawner(BeakerhubImageSpawner): + """Launch a Beaker session as an ECS runtime process.""" + + runtime = Instance(AwsEcsRuntime, allow_none=False) + + cluster_name = Unicode(config=True, help="ECS cluster name or ARN.") + task_definition = Unicode( + config=True, + help="Optional ECS task-definition family, revision, or ARN. When unset, " + "BeakerHub registers and reuses an image-specific definition.", + ) + container_name = Unicode( + "task-container", + config=True, + help="Name of the notebook container in the ECS task definition.", + ) + container_tags = Dict( + Unicode(), Unicode(), default_value={}, config=True, + help="Tags applied to every ECS notebook task.", + ) + subnets = List( + Unicode(), config=True, + help="Subnets used by the task's awsvpc network configuration.", + ) + security_groups = List(Unicode(), config=True) + assign_public_ip = Bool(config=True) + launch_options = Dict( + config=True, + help="ECS run_task options for notebook tasks.", + ) + efs_volume_configuration = Dict( + default_value={}, + config=True, + help="Optional ECS efsVolumeConfiguration mapping for the notebook task.", + ) + efs_mount_path = Unicode( + "", + config=True, + help="Container path for the optional EFS volume mount.", + ) + efs_volume_name = Unicode( + "beakerhub-efs", + config=True, + help="Task-definition volume name for the optional EFS mount.", + ) + task_overrides = Dict(default_value=None, allow_none=True, config=True) + log_group = Unicode(config=True) + log_stream_prefix = Unicode(config=True) + execution_role_arn = Unicode(config=True) + task_role_arn = Unicode(config=True) + cpu_architecture: Literal["X86_64", "ARM64"] = traitlets.Enum( + ["X86_64", "ARM64"], config=True + ) + task_definition_name = Unicode(config=True) + task_group = Unicode(config=True) + launch_type: Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"] = ( + traitlets.Enum( + ["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], config=True + ) + ) + default_task_cpu = Unicode(config=True) + default_task_memory = Unicode(config=True) + default_task_ephemeral_storage = traitlets.Int(config=True) + + task_arn: str | None = None + + @default("http_timeout") + def _default_http_timeout(self): + return 600 + + @default("runtime") + def _default_runtime(self) -> AwsEcsRuntime: + app = self.user.settings.get("app") + runtime = getattr(app, "runtime", None) + if not isinstance(runtime, AwsEcsRuntime): + raise RuntimeError( + "BeakerAwsECSSpawner requires the application to use AwsEcsRuntime" + ) + return runtime + + @default("cluster_name") + def _default_cluster_name(self) -> str: + return self.runtime.cluster_name + + @default("subnets") + def _default_subnets(self) -> list[str]: + return list(self.runtime.subnets) + + @default("security_groups") + def _default_security_groups(self) -> list[str]: + return list(self.runtime.security_groups) + + @default("assign_public_ip") + def _default_assign_public_ip(self) -> bool: + return self.runtime.assign_public_ip + + @default("launch_options") + def _default_launch_options(self) -> dict[str, Any]: + return dict(self.runtime.launch_options) + + @default("task_overrides") + def _default_task_overrides(self) -> dict[str, Any] | None: + return ( + dict(self.runtime.task_overrides) + if self.runtime.task_overrides is not None + else None + ) + + @default("log_group") + def _default_log_group(self) -> str: + return self.runtime.log_group + + @default("log_stream_prefix") + def _default_log_stream_prefix(self) -> str: + return self.runtime.log_stream_prefix + + @default("execution_role_arn") + def _default_execution_role_arn(self) -> str: + return self.runtime.execution_role_arn + + @default("task_role_arn") + def _default_task_role_arn(self) -> str: + return self.runtime.task_role_arn + + @default("cpu_architecture") + def _default_cpu_architecture(self) -> Literal["X86_64", "ARM64"]: + return self.runtime.cpu_architecture + + @default("task_definition_name") + def _default_task_definition_name(self) -> str: + identifier = b"\x1f".join( + str(value).encode("utf-8") for value in (self.user.id, self.image) + ) + return f"beakerhub-notebook_{hashlib.sha256(identifier).hexdigest()}" + + @default("task_group") + def _default_task_group(self) -> str: + return self.runtime.task_group + + @default("launch_type") + def _default_launch_type( + self, + ) -> Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"]: + return self.runtime.launch_type + + @default("default_task_cpu") + def _default_task_cpu(self) -> str: + return self.runtime.default_task_cpu + + @default("default_task_memory") + def _default_task_memory(self) -> str: + return self.runtime.default_task_memory + + @default("default_task_ephemeral_storage") + def _default_task_ephemeral_storage(self) -> int: + return self.runtime.default_task_ephemeral_storage + + @property + def volume_configs(self) -> list[dict[str, Any]]: + """Return ECS managed-volume configuration overrides, if configured.""" + return [] + + def _task_tags(self) -> dict[str, str]: + tags = dict(self.container_tags) + # The session tag is owned by BeakerHub and must not be overridden. + tags["beaker-session"] = self.session_id + return tags + + def _efs_volumes(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Build optional EFS task-definition volumes and container mounts.""" + configured = bool(self.efs_volume_configuration) + has_mount_path = bool(self.efs_mount_path) + if configured != has_mount_path: + raise ValueError( + "efs_volume_configuration and efs_mount_path must be configured together" + ) + if not configured: + return [], [] + + file_system_id = self.efs_volume_configuration.get("fileSystemId") + if not file_system_id: + raise ValueError("efs_volume_configuration.fileSystemId must be configured") + + user_subdir = f"/user-storage/{strip_and_hash(self.user.name)}" + access_point_id = self._efs_access_point_id( + file_system_id, + str(self.user.id), + user_subdir, + ) + volume_configuration = dict(self.efs_volume_configuration) + authorization_config = dict(volume_configuration.get("authorizationConfig", {})) + authorization_config["accessPointId"] = access_point_id + volume_configuration["authorizationConfig"] = authorization_config + volume_configuration.pop("rootDirectory", None) + + return ( + [ + { + "name": self.efs_volume_name, + "efsVolumeConfiguration": volume_configuration, + } + ], + [ + { + "sourceVolume": self.efs_volume_name, + "containerPath": self.efs_mount_path, + "readOnly": False, + } + ], + ) + + def _efs_access_point_id( + self, + file_system_id: str, + user_id: str, + user_subdir: str, + ) -> str: + """Find or create the EFS access point for one user's storage.""" + efs_client = boto3.client("efs", region_name=self.runtime.aws_region) + tagging_client = boto3.client( + "resourcegroupstaggingapi", + region_name=self.runtime.aws_region, + ) + access_point_id = self._find_efs_access_point( + efs_client, + tagging_client, + file_system_id, + user_id, + user_subdir, + ) + if access_point_id is not None: + return access_point_id + + response = efs_client.create_access_point( + ClientToken=self._efs_access_point_token(file_system_id, user_id), + FileSystemId=file_system_id, + RootDirectory={ + "Path": user_subdir, + "CreationInfo": { + "OwnerUid": 1000, + "OwnerGid": 1000, + "Permissions": "755", + }, + }, + Tags=[ + {"Key": "beakerhub-user-id", "Value": user_id}, + {"Key": "beakerhub-user-name", "Value": str(self.user.name)}, + {"Key": "beakerhub-user-subdir", "Value": user_subdir}, + ], + ) + access_point_id = response["AccessPointId"] + self._wait_for_efs_access_point(efs_client, access_point_id) + return access_point_id + + @staticmethod + def _efs_access_point_token(file_system_id: str, user_id: str) -> str: + return hashlib.sha256( + b"\x1f".join( + value.encode("utf-8") for value in (file_system_id, user_id) + ) + ).hexdigest() + + @staticmethod + def _find_efs_access_point( + efs_client: Any, + tagging_client: Any, + file_system_id: str, + user_id: str, + user_subdir: str, + ) -> str | None: + """Find the matching access point among resources tagged for this user.""" + pagination_token = "" + while True: + request: dict[str, Any] = { + "TagFilters": [{"Key": "beakerhub-user-id", "Values": [user_id]}], + "ResourceTypeFilters": ["elasticfilesystem:access-point"], + } + if pagination_token: + request["PaginationToken"] = pagination_token + response = tagging_client.get_resources(**request) + for resource in response.get("ResourceTagMappingList", []): + candidate_id = resource["ResourceARN"].rsplit("/", maxsplit=1)[-1] + described = efs_client.describe_access_points( + AccessPointId=candidate_id + ) + access_points = described.get("AccessPoints", []) + if not access_points: + continue + access_point = access_points[0] + if ( + access_point.get("FileSystemId") == file_system_id + and access_point.get("RootDirectory", {}).get("Path") == user_subdir + ): + return access_point["AccessPointId"] + if not response.get("PaginationToken"): + return None + pagination_token = response["PaginationToken"] + + @staticmethod + def _wait_for_efs_access_point(efs_client: Any, access_point_id: str) -> None: + """Wait until a newly created EFS access point is usable.""" + for _ in range(60): + described = efs_client.describe_access_points(AccessPointId=access_point_id) + access_points = described.get("AccessPoints", []) + state = access_points[0].get("LifeCycleState") if access_points else None + if state == "available": + return + if state in {"deleted", "deleting", "error"}: + raise RuntimeError( + f"EFS access point {access_point_id} entered state {state!r}" + ) + time.sleep(1) + raise RuntimeError( + f"EFS access point {access_point_id} did not become available" + ) + + def get_env(self) -> dict[str, str]: + """Add the hub URLs needed by an ECS-hosted notebook server.""" + env = super().get_env() + hub_host = os.environ.get("HOSTNAME", "hub") + env.update( + { + "JUPYTERHUB_API_URL": f"http://{hub_host}:8888/api", + "JUPYTERHUB_ACTIVITY_URL": ( + f"http://{hub_host}:8888/api/users/matt@jataware.com/activity" + ), + "JUPYTERHUB_SERVICE_URL": f"http://{hub_host}:8888/", + } + ) + return env + + def _task_overrides(self) -> dict[str, Any]: + """Merge session environment into runtime task overrides. + + Session values are run-task overrides, rather than task-definition + values, so image-specific definitions can be reused by multiple users. + """ + overrides = dict(self.task_overrides or {}) + containers: list[dict[str, Any]] = [] + session_container: dict[str, Any] = {"name": self.container_name} + environment: dict[str, str] = {} + for container in overrides.pop("containerOverrides", []): + container = dict(container) + if container.get("name") != self.container_name: + containers.append(container) + continue + environment.update( + {item["name"]: item["value"] for item in container.pop("environment", [])} + ) + session_container.update(container) + environment.update({name: str(value) for name, value in self.get_env().items()}) + if environment: + session_container["environment"] = [ + {"name": name, "value": value} for name, value in environment.items() + ] + containers.append(session_container) + overrides["containerOverrides"] = containers + return overrides + + def _config_writer_container(self) -> dict[str, Any]: + """Build the sidecar that writes the generated notebook config.""" + return { + "name": _CONFIG_WRITER_NAME, + "image": _CONFIG_WRITER_IMAGE, + "command": [ + "/bin/sh", + "-ec", + "printf '%s' \"$BEAKER_NOTEBOOK_CONFIG\" > " + f"{_CONFIG_WRITER_MOUNT_PATH}/beaker_config.py", + ], + "environment": [ + { + "name": "BEAKER_NOTEBOOK_CONFIG", + "value": str(self._notebook_config()), + } + ], + "mountPoints": [ + { + "containerPath": _CONFIG_WRITER_MOUNT_PATH, + "sourceVolume": _CONFIG_VOLUME_NAME, + } + ], + "essential": False, + } + + def _definition(self) -> AwsEcsDefinition: + """Build the ECS definition for this notebook session.""" + if not self.container_name: + raise ValueError("BeakerAwsECSSpawner.container_name must be configured") + + launch_options = dict(self.launch_options) + if self.volume_configs: + launch_options["volumeConfigurations"] = self.volume_configs + volumes, mount_points = self._efs_volumes() + + volumes.append({"name": _CONFIG_VOLUME_NAME}) + mount_points.append( + { + "containerPath": _NOTEBOOK_CONFIG_PATH, + "sourceVolume": _CONFIG_VOLUME_NAME, + } + ) + + # This flag allows an AWS user/accout with a proper role to "log into" running tasks. + # As such, it is only enabled when running in debug mode. + launch_options["enableExecuteCommand"] = self.debug + + return AwsEcsDefinition( + runtime=self.runtime, + image=self.image, + task_definition_arn=self.task_definition or None, + container_name=self.container_name, + tags=self._task_tags(), + definition_tags={ + "beakerhub-user-id": str(self.user.id), + "beakerhub-image-uri": self.image, + }, + cluster_name=self.cluster_name, + subnets=list(self.subnets), + security_groups=list(self.security_groups), + assign_public_ip=self.assign_public_ip, + log_group=self.log_group, + log_stream_prefix=self.log_stream_prefix, + execution_role_arn=self.execution_role_arn, + task_role_arn=self.task_role_arn, + cpu_architecture=self.cpu_architecture, + task_definition_name=self.task_definition_name, + task_group=self.task_group, + launch_type=self.launch_type, + launch_options=launch_options, + task_overrides=self._task_overrides(), + cpu=self.default_task_cpu, + memory=self.default_task_memory, + ephemeral_storage=self.default_task_ephemeral_storage, + volumes=volumes, + mount_points=mount_points, + sidecar_containers=[self._config_writer_container()], + ) + + def _process(self) -> AwsEcsProcess: + """Reconstruct the runtime process for this session, if it was launched.""" + return AwsEcsProcess( + self._definition(), + runtime=self.runtime, + external_id=self.task_arn, + process_type="service", + ) + + @cache + def _notebook_config(self) -> str: + from beakerhub.app import BeakerHub + beakerhub = BeakerHub.instance() + subdomain_host = beakerhub.subdomain_host + + c = Config() + c.BaseBeakerApp.allow_origin_pat = rf"^{subdomain_host}$" + c.BaseBeakerApp.tornado_settings = { + "headers": { + "Content-Security-Policy": f"frame-ancestors 'self' {subdomain_host}" + } + } + c.BaseBeakerApp.identity_provider_class = ( + "beakerhub.auth.node.BeakerhubNodeIdentityProvider" + ) + c.BaseBeakerApp.authorizer_class = ( + "beakerhub.auth.node.BeakerhubNodeAuthorizer" + ) + c.BaseBeakerApp.secrets_manager_class = ( + "beakerhub.services.secrets.beakerhub.BeakerhubSecretsManager" + ) + c.BeakerhubSecretsManager.policy_override_env_key_suffix = "_secret_policy" + c.FileNotebookManager.notebook_path = ".notebooks" + c.FileNotebookManager.snapshot_path = ".notebooks" + + lines = [ + "# Beaker Notebook Service Configuration File", + "", + "c = get_config() # noqa # type: ignore", + "", + ] + + for cls_name, cls_configs in c.items(): + for config_name, config_value in cls_configs.items(): + lines.append(f"c.{cls_name}.{config_name} = {repr(config_value)}") + config_text = "\n".join(lines) + + return config_text + + async def start(self) -> tuple[str, int]: + """Launch the session task and retain its ARN for later lifecycle calls.""" + process = AwsEcsProcess.start( + self._definition(), + runtime=self.runtime, + process_type="service", + ) + self.task_arn = process.external_id + local_ip = await self._wait_for_task_ip(process=process) + + port = 8888 + self.log.warning("ECS notebook task is available at http://%s:%s", local_ip, port) + return str(local_ip), port + + async def stop(self, now: bool = False) -> None: + """Stop the session task and wait for ECS to deprovision it.""" + if self.task_arn: + process = self._process() + process.stop() + await process.await_completion() + + async def poll(self) -> int | None: + """Map normalized ECS process status to JupyterHub poll semantics.""" + if not self.task_arn: + return 0 + status: ProcessStatus = self._process().status + if status.state in {"pending", "running"}: + return None + if status.exit_code is not None: + return status.exit_code + return 0 if status.state == "completed" else 1 + + def get_state(self) -> dict[str, Any]: + state = super().get_state() + if self.task_arn: + state["task_arn"] = self.task_arn + return state + + def load_state(self, state: dict[str, Any]) -> None: + super().load_state(state) + self.task_arn = state.get("task_arn") + + def clear_state(self) -> None: + super().clear_state() + self.task_arn = None + + async def _wait_for_task_ip( + self, + process: AwsEcsProcess, + timeout: float = 120.0, + ) -> str: + """Return the task's private IPv4 address once ECS has assigned one.""" + def _request_ip() -> str | None: + task = self.runtime.describe_task( + process.definition.cluster_name, + process.external_id, + ) + if task is None: + raise RuntimeError(f"Task `{process.external_id}` was not found") + if task.get("lastStatus") == "STOPPED": + container = next( + ( + item + for item in task.get("containers", []) + if item.get("name") == process.definition.container_name + ), + None, + ) + reason = ( + (container or {}).get("reason") + or task.get("stoppedReason") + or "Unknown failure" + ) + raise RuntimeError( + f"Task `{process.external_id}` stopped before receiving an " + f"internal IP: {reason}" + ) + return next( + ( + detail["value"] + for attachment in task.get("attachments", []) + if attachment.get("type") == "ElasticNetworkInterface" + for detail in attachment.get("details", []) + if detail.get("name") == "privateIPv4Address" + ), + None, + ) + + start = time.monotonic() + while not (task_ip := _request_ip()): + if time.monotonic() - start > timeout: + raise TimeoutError( + f"Task `{process.external_id}` did not receive an internal IP " + f"within {timeout} seconds." + ) + await asyncio.sleep(5) + return task_ip diff --git a/src/beakerhub/spawner/kubernetes.py b/src/beakerhub/services/spawner/base.py similarity index 70% rename from src/beakerhub/spawner/kubernetes.py rename to src/beakerhub/services/spawner/base.py index 1bce24f..33c0a29 100644 --- a/src/beakerhub/spawner/kubernetes.py +++ b/src/beakerhub/services/spawner/base.py @@ -1,22 +1,14 @@ -import json -from typing import cast +from typing import Any, cast +from uuid import uuid4 from jupyterhub.spawner import Spawner -from kubespawner.spawner import KubeSpawner -from traitlets import default, validate, Unicode, Dict, List -from traitlets.config import Application +from traitlets import Dict, Unicode, default, validate +from beakerhub import orm from beakerhub.auth.user import BeakerhubUser -from beakerhub.spawner.base import BeakerSpawner from beakerhub.services.secrets import VAULT_ENV_VAR_LIST_KEY -from beakerhub import orm - - -class BeakerKubeSpawner(KubeSpawner, BeakerSpawner): - default_registry = Unicode().tag(config=True) - default_image = Unicode().tag(config=True) - default_tag = Unicode(default_value="latest").tag(config=True) +class BeakerSpawner(Spawner): default_beaker_context = Unicode( default_value="default", @@ -37,63 +29,38 @@ class BeakerKubeSpawner(KubeSpawner, BeakerSpawner): "Resolved from the vault at spawn time and passed to the node as JSON." ) - @staticmethod - def image_has_defined_registry(image: str) -> bool: - image_parts = image.split("/") - if len(image_parts) == 1: - return False - return "." in image_parts[0] or ":" in image_parts[0] - - @validate("default_registry") - def _validate_default_registry(self, proposal): - return proposal["value"].rstrip("/") - - @default("delete_stopped_pods") - def _default_delete_stopped_pods(self): - return False - - @default("namespace") - def _default_namespace(self): - return "beakerhub" - - @default("default_image") - def _default_default_image(self): - return f"{self.default_registry}/beakerhub/default-node:{self.default_tag}" - - @validate("default_image") - def _validate_default_image(self, proposal): - value = proposal["value"] - if not self.image_has_defined_registry(value): - return f"{self.default_registry}/{value}" - else: - return value - @default("image") - def _default_image(self): - return self.default_image + def __init__(self, **kwargs: Any) -> None: + domain = kwargs.pop("domain", None) + super().__init__(**kwargs) + if domain: + self.proxy_spec = f"{self.name}.{domain}/" - @validate("image") - def _validate_image(self, proposal): - value = proposal["value"] - if not self.image_has_defined_registry(value): - return f"{self.default_registry}/{value}" - else: - return value + @property + def name(self) -> str: + name = super().name + return name - @default("pod_name_template") - def _default_pod_name_template(self): - return "session-{user_server}" + @name.setter + def name(self, value: str) -> str: + if self.orm_spawner and self.orm_spawner.name != value: + self.orm_spawner.name = value + self.orm_spawner.save() + return value @property - def node_type(self) -> str | None: - pod_name = self.image.split("/")[-1] - if "node" not in pod_name: - return None - return pod_name + def session_id(self) -> str: + if not self.name: + self.name = str(uuid4()) + return self.name def get_env(self): - user: BeakerhubUser = cast(BeakerhubUser, self.user) env = super().get_env() + return self._extend_env(env) + + def _extend_env(self, env: dict[str, str]) -> dict[str, str]: + """Add BeakerHub-specific values to a runtime environment.""" + user = cast(BeakerhubUser, self.user) env.setdefault("JUPYTER_BASE_URL", user.server_url(server_name=self.name)) env.setdefault("BEAKER_DEFAULT_CONTEXT", self.beaker_context or self.default_beaker_context) env.setdefault("BEAKERHUB_USER", user.name) @@ -114,35 +81,42 @@ def get_env(self): return env - @staticmethod - def apply_user_options(spawner: "BeakerKubeSpawner", user_options: dict): + def start(self): + raise NotImplementedError() + + def stop(self, now=False): + raise NotImplementedError() + + def poll(self): + raise NotImplementedError() + + def apply_user_options(self, _spawner, user_options: dict): node_record: orm.NodeImages | None = None context_slug: str | None = None if (node_slug := user_options.get("nodeSlug", None)): - node_record = spawner.db.query(orm.NodeImages).filter(orm.NodeImages.slug == node_slug).first() - if node_record: - tag = ("debug" if spawner.debug else node_record.default_tag) or spawner.default_tag - image_spec = f"{node_record.default_registry}/{node_record.repository}:{tag}" - spawner.image = image_spec + node_record = self.db.query(orm.NodeImages).filter(orm.NodeImages.slug == node_slug).first() if (context := user_options.get("contextSlug", None)): if ':' in context: context = context.split(":")[-1] context_slug = context - spawner.beaker_context = context + self.beaker_context = context + + if (context_config := user_options.get("contextOptions", None)): + self.context_config = context_config # Inject secrets from the vault: globals first, then node-specific overrides # Build the set of secret IDs explicitly disabled for this context disabled_secret_ids: set[int] = set() if context_slug: context_record = ( - spawner.db.query(orm.Context) + self.db.query(orm.Context) .filter(orm.Context.slug == context_slug) .first() ) if context_record: - disabled_rows = spawner.db.execute( + disabled_rows = self.db.execute( orm.beaker_context_secrets.select().where( orm.beaker_context_secrets.c.context_id == context_record.id, orm.beaker_context_secrets.c.enabled == False, @@ -172,7 +146,7 @@ def collect(secret: orm.NodeSecret) -> None: # Global secrets (node_image_id IS NULL) global_secrets = ( - spawner.db.query(orm.NodeSecret) + self.db.query(orm.NodeSecret) .filter(orm.NodeSecret.node_image_id.is_(None)) .all() ) @@ -182,26 +156,80 @@ def collect(secret: orm.NodeSecret) -> None: # Node-specific secrets (override globals) if node_record: node_secrets = ( - spawner.db.query(orm.NodeSecret) + self.db.query(orm.NodeSecret) .filter(orm.NodeSecret.node_image_id == node_record.id) .all() ) for secret in node_secrets: collect(secret) - spawner.node_env = secrets_env or {} - spawner.node_policy_overrides = secrets_policies or {} + self.node_env = secrets_env or {} + self.node_policy_overrides = secrets_policies or {} - # Keep K8s secretRef as fallback for secrets not yet migrated to the vault - if node_slug: - secret_name = f"{node_slug}-secrets" - env_from: list = spawner.extra_container_config.setdefault("envFrom", []) - env_from.append({ - "secretRef": { - "name": secret_name, - "optional": True # Don't fail if secret doesn't exist - } - }) +class BeakerhubImageSpawner(BeakerSpawner): - if (context_config := user_options.get("contextOptions", None)): - spawner.context_config = context_config + default_registry = Unicode( + config=True, + ).tag(config=True) + default_image = Unicode( + config=True, + ).tag(config=True) + default_tag = Unicode( + "latest", + config=True, + ) + + image = Unicode( + "beakerhub/default-node:latest", + config=True, + help=""" + Docker image to use for spawning user's containers. + """, + ) + + @staticmethod + def image_has_defined_registry(image: str) -> bool: + image_parts = image.split("/") + if len(image_parts) == 1: + return False + return "." in image_parts[0] or ":" in image_parts[0] + + @validate("default_registry") + def _validate_default_registry(self, proposal): + return proposal["value"].rstrip("/") + + @default("default_image") + def _default_default_image(self): + return f"{self.default_registry}/beakerhub/default-node:{self.default_tag}" + + @validate("default_image") + def _validate_default_image(self, proposal): + value = proposal["value"] + if not self.image_has_defined_registry(value): + return f"{self.default_registry}/{value}" + else: + return value + + @default("image") + def _default_image(self): + return self.default_image + + @validate("image") + def _validate_image(self, proposal): + value = proposal["value"] + if not self.image_has_defined_registry(value): + return f"{self.default_registry}/{value}" + else: + return value + + def apply_user_options(self, _spawner, user_options: dict): + node_record: orm.NodeImages | None = None + + super().apply_user_options(_spawner, user_options) + + if (node_slug := user_options.get("nodeSlug", None)): + node_record = self.db.query(orm.NodeImages).filter(orm.NodeImages.slug == node_slug).first() + if node_record: + tag = ("debug" if self.debug else node_record.default_tag) or self.default_tag + image_spec = f"{node_record.default_registry}/{node_record.repository}:{tag}" + self.image = image_spec diff --git a/src/beakerhub/services/spawner/kubernetes_spawner.py b/src/beakerhub/services/spawner/kubernetes_spawner.py new file mode 100644 index 0000000..45bad96 --- /dev/null +++ b/src/beakerhub/services/spawner/kubernetes_spawner.py @@ -0,0 +1,53 @@ +import json +from typing import cast + +from jupyterhub.spawner import Spawner +from kubespawner.spawner import KubeSpawner +from traitlets import default, validate, Unicode, Dict, List +from traitlets.config import Application + +from beakerhub.services.spawner.base import BeakerSpawner, BeakerhubImageSpawner +from beakerhub import orm + + +class BeakerKubeSpawner(KubeSpawner, BeakerhubImageSpawner): + + image = KubeSpawner.image + + def get_env(self): + """Add BeakerHub variables after KubeSpawner builds its environment.""" + return BeakerSpawner._extend_env(self, super().get_env()) + + @default("delete_stopped_pods") + def _default_delete_stopped_pods(self): + return False + + @default("namespace") + def _default_namespace(self): + return "beakerhub" + + @default("pod_name_template") + def _default_pod_name_template(self): + return "session-{user_server}" + + @property + def node_type(self) -> str | None: + pod_name = self.image.split("/")[-1] + if "node" not in pod_name: + return None + return pod_name + + def apply_user_options(self, _spawner, user_options: dict): + # Equivilent to calling method via super() when multiple inheritence + BeakerhubImageSpawner.apply_user_options(self, _spawner, user_options) + + # Keep K8s secretRef as fallback for secrets not yet migrated to the vault + if (node_slug := user_options.get("nodeSlug", None)): + secret_name = f"{node_slug}-secrets" + env_from: list = self.extra_container_config.setdefault("envFrom", []) + env_from.append({ + "secretRef": { + "name": secret_name, + "optional": True # Don't fail if secret doesn't exist + } + }) diff --git a/src/beakerhub/spawner/provisioner.py b/src/beakerhub/services/spawner/provisioner.py similarity index 61% rename from src/beakerhub/spawner/provisioner.py rename to src/beakerhub/services/spawner/provisioner.py index 8deb26f..187fb2d 100644 --- a/src/beakerhub/spawner/provisioner.py +++ b/src/beakerhub/services/spawner/provisioner.py @@ -2,9 +2,9 @@ from beaker_notebook.services.kernel.provisioner import BeakerLocalProvisioner -class KubernetesLocalProvisioner(BeakerLocalProvisioner): +class BeakerhubLocalProvisioner(BeakerLocalProvisioner): """ - Custom Provisioner that is ensures that launched kernel logs are forwarded to the main stdout/stderr for proper logging in kubernetes + Forward launched kernel logs to the parent process output streams. """ async def launch_kernel(self, cmd, **kwargs): kwargs.setdefault("stdout", sys.stdout) diff --git a/src/beakerhub/services/task/__init__.py b/src/beakerhub/services/task/__init__.py new file mode 100644 index 0000000..aa5dc76 --- /dev/null +++ b/src/beakerhub/services/task/__init__.py @@ -0,0 +1,5 @@ +"""Background task-runner services.""" + +from .base import BaseTaskRunnerService + +__all__ = ["BaseTaskRunnerService"] diff --git a/src/beakerhub/services/task/aws_ecs_task_runner.py b/src/beakerhub/services/task/aws_ecs_task_runner.py new file mode 100644 index 0000000..e4fc9ec --- /dev/null +++ b/src/beakerhub/services/task/aws_ecs_task_runner.py @@ -0,0 +1,168 @@ +"""AWS ECS implementation of the task-runner service.""" + +from typing import Any, Literal + +import traitlets + +from beakerhub.runtimes.aws_ecs import AwsEcsDefinition, AwsEcsProcess, AwsEcsRuntime +from beakerhub.services.task.base import BaseTaskRunnerService +from beakerhub.tasks.base import BaseImageTask, BaseTaskDefinition + + +class AwsEcsTaskRunnerService(BaseTaskRunnerService): + """Adapt image tasks to ECS runtime processes.""" + + runtime = traitlets.Instance(AwsEcsRuntime, allow_none=False) + + log_group = traitlets.Unicode(config=True) + log_stream_prefix = traitlets.Unicode(config=True) + execution_role_arn = traitlets.Unicode(config=True) + task_role_arn = traitlets.Unicode(config=True) + subnets = traitlets.List(traitlets.Unicode(), config=True) + security_groups = traitlets.List(traitlets.Unicode(), config=True) + assign_public_ip = traitlets.Bool(config=True) + cpu_architecture: Literal["X86_64", "ARM64"] = traitlets.Enum( + ["X86_64", "ARM64"], config=True + ) + task_definition_name = traitlets.Unicode(config=True) + cluster_name = traitlets.Unicode(config=True) + task_group = traitlets.Unicode(config=True) + launch_type: Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"] = ( + traitlets.Enum( + ["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], config=True + ) + ) + launch_options = traitlets.Dict(config=True) + task_overrides = traitlets.Dict(default_value=None, allow_none=True, config=True) + default_task_cpu = traitlets.Unicode(config=True) + default_task_memory = traitlets.Unicode(config=True) + default_task_ephemeral_storage = traitlets.Int(config=True) + + @traitlets.default("runtime") + def _default_runtime(self) -> AwsEcsRuntime: + parent_runtime = getattr(self.parent, "runtime", None) + if isinstance(parent_runtime, AwsEcsRuntime): + return parent_runtime + return AwsEcsRuntime(parent=self) + + @traitlets.default("log_group") + def _default_log_group(self) -> str: + return self.runtime.log_group + + @traitlets.default("log_stream_prefix") + def _default_log_stream_prefix(self) -> str: + return self.runtime.log_stream_prefix + + @traitlets.default("execution_role_arn") + def _default_execution_role_arn(self) -> str: + return self.runtime.execution_role_arn + + @traitlets.default("task_role_arn") + def _default_task_role_arn(self) -> str: + return self.runtime.task_role_arn + + @traitlets.default("subnets") + def _default_subnets(self) -> list[str]: + return list(self.runtime.subnets) + + @traitlets.default("security_groups") + def _default_security_groups(self) -> list[str]: + return list(self.runtime.security_groups) + + @traitlets.default("assign_public_ip") + def _default_assign_public_ip(self) -> bool: + return self.runtime.assign_public_ip + + @traitlets.default("cpu_architecture") + def _default_cpu_architecture(self) -> Literal["X86_64", "ARM64"]: + return self.runtime.cpu_architecture + + @traitlets.default("task_definition_name") + def _default_task_definition_name(self) -> str: + return self.runtime.task_definition_name + + @traitlets.default("cluster_name") + def _default_cluster_name(self) -> str: + return self.runtime.cluster_name + + @traitlets.default("task_group") + def _default_task_group(self) -> str: + return self.runtime.task_group + + @traitlets.default("launch_type") + def _default_launch_type( + self, + ) -> Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"]: + return self.runtime.launch_type + + @traitlets.default("launch_options") + def _default_launch_options(self) -> dict[str, Any]: + return dict(self.runtime.launch_options) + + @traitlets.default("task_overrides") + def _default_task_overrides(self) -> dict[str, Any] | None: + return ( + dict(self.runtime.task_overrides) + if self.runtime.task_overrides is not None + else None + ) + + @traitlets.default("default_task_cpu") + def _default_task_cpu(self) -> str: + return self.runtime.default_task_cpu + + @traitlets.default("default_task_memory") + def _default_task_memory(self) -> str: + return self.runtime.default_task_memory + + @traitlets.default("default_task_ephemeral_storage") + def _default_task_ephemeral_storage(self) -> int: + return self.runtime.default_task_ephemeral_storage + + def submit(self, task: BaseTaskDefinition) -> AwsEcsProcess: + """Create an ECS definition and launch one image task.""" + if not isinstance(task, BaseImageTask): + raise ValueError( + "AwsEcsTaskRunnerService only supports image tasks, not " + f"{task.task_type!r}" + ) + return AwsEcsProcess.start(self._definition(task), runtime=self.runtime) + + def _definition(self, task: BaseImageTask | None = None) -> AwsEcsDefinition: + """Translate task and service overrides to an ECS runtime definition.""" + return AwsEcsDefinition( + runtime=self.runtime, + image=task.image if task else "", + entrypoint=list(task.entrypoint) if task else [], + command=list(task.command) if task else [], + working_directory=task.working_directory if task else None, + environment=dict(task.environment) if task else {}, + tags={"beakerhub/task-type": task.task_type} if task else {}, + cpu=self.default_task_cpu, + memory=self.default_task_memory, + ephemeral_storage=self.default_task_ephemeral_storage, + log_group=self.log_group, + log_stream_prefix=self.log_stream_prefix, + execution_role_arn=self.execution_role_arn, + task_role_arn=self.task_role_arn, + subnets=list(self.subnets), + security_groups=list(self.security_groups), + assign_public_ip=self.assign_public_ip, + cpu_architecture=self.cpu_architecture, + task_definition_name=self.task_definition_name, + cluster_name=self.cluster_name, + task_group=self.task_group, + launch_type=self.launch_type, + launch_options=dict(self.launch_options), + task_overrides=( + dict(self.task_overrides) if self.task_overrides is not None else None + ), + ) + + def get_process(self, external_id: str) -> AwsEcsProcess: + """Reconstruct an ECS process from its persisted task ARN.""" + return AwsEcsProcess( + self._definition(), + runtime=self.runtime, + external_id=external_id, + ) diff --git a/src/beakerhub/services/task/base.py b/src/beakerhub/services/task/base.py new file mode 100644 index 0000000..5aee026 --- /dev/null +++ b/src/beakerhub/services/task/base.py @@ -0,0 +1,53 @@ +"""Task-runner adapters for persisted BeakerHub background tasks.""" + +from typing import TYPE_CHECKING, TypeAlias + +from beakerhub.runtimes.base import ( + BaseProcess, + ProcessOutput, + ProcessState, + ProcessStatus, +) +from traitlets.config import LoggingConfigurable + +if TYPE_CHECKING: + from beakerhub.tasks.base import BaseTaskDefinition + + +# Keep task-service imports stable while using the runtime lifecycle contract. +TaskState: TypeAlias = ProcessState +TaskStatus = ProcessStatus +TaskOutput = ProcessOutput + + +class BaseTaskRunnerService(LoggingConfigurable): + """Adapt logical BeakerHub tasks to provider runtime processes. + + Persisted tasks retain only an external provider identifier. Provider-specific + subclasses reconstruct a process handle from that identifier when the task + service later polls status, reads output, or requests cleanup. + """ + + def submit(self, task: "BaseTaskDefinition") -> BaseProcess: + """Launch a runtime process for a logical task.""" + raise NotImplementedError + + def get_process(self, external_id: str) -> BaseProcess: + """Reconstruct a process handle for a persisted external identifier.""" + raise NotImplementedError + + def get_status(self, external_id: str) -> ProcessStatus: + """Return the current status of a persisted runtime process.""" + return self.get_process(external_id).status + + def get_output(self, external_id: str) -> ProcessOutput | None: + """Return retained output for a persisted runtime process.""" + return self.get_process(external_id).collect_output() + + def delete(self, external_id: str) -> None: + """Request cleanup of a persisted runtime process.""" + self.get_process(external_id).stop() + + def reap_stale_tasks(self) -> None: + """Find and clean backend workloads orphaned from normal completion.""" + raise NotImplementedError diff --git a/src/beakerhub/services/task/handlers.py b/src/beakerhub/services/task/handlers.py new file mode 100644 index 0000000..06f28f8 --- /dev/null +++ b/src/beakerhub/services/task/handlers.py @@ -0,0 +1,111 @@ +"""API handlers for background tasks.""" + +import json +import logging +from datetime import datetime, timezone +from typing import Any + +from jupyterhub.apihandlers import APIHandler +from jupyterhub.scopes import needs_scope +from tornado import web + +from beakerhub.orm import NodeImages, NodeImageTask +from beakerhub.tasks.image_import.task import ImageImportTask, launch_import_task + + +log = logging.getLogger(__name__) + + +class NodeImageImportHandler(APIHandler): + """Trigger a context import for a node image.""" + + @needs_scope("admin:users") + async def post(self, image_id: str): + node = self.db.query(NodeImages).filter(NodeImages.id == int(image_id)).first() + if not node: + raise web.HTTPError(404, f"Node image not found: {image_id}") + + try: + task = launch_import_task(self.db, self.settings.get("app"), node) + except ValueError as error: + raise web.HTTPError(409, str(error)) from error + except RuntimeError as error: + raise web.HTTPError(500, str(error)) from error + + self.set_header("Content-Type", "application/json") + self.write(json.dumps({ + "task_id": task.id, + "job_name": task.job_name, + "status": task.status, + })) + + +class NodeImageImportStatusHandler(APIHandler): + """Poll and collect the result of a node-image import task.""" + + def compute_etag(self) -> None: + return None + + @needs_scope("admin:users") + async def get(self, image_id: str): + node = self.db.query(NodeImages).filter(NodeImages.id == int(image_id)).first() + if not node: + raise web.HTTPError(404, f"Node image not found: {image_id}") + + task = self.db.query(NodeImageTask).filter( + NodeImageTask.node_image_id == node.id, + NodeImageTask.task_type == "context_import", + ).order_by(NodeImageTask.created_at.desc()).first() + + if not task: + self.set_header("Content-Type", "application/json") + self.write(json.dumps({"status": "none", "message": "No import has been run"})) + return + + if task.status == "running" and task.job_name: + try: + runner = self.settings["app"].task_runner + status = runner.get_status(task.job_name) + definition = ImageImportTask.from_node_image(node) + if status.state == "completed": + output = runner.get_output(task.job_name) + if output is None: + raise RuntimeError("Task completed but no output is available") + definition.on_success(task, status, output) + task.status = "completed" + task.error = None + task.updated_at = datetime.now(timezone.utc) + self.db.commit() + runner.delete(task.job_name) + elif status.state == "failed": + output = runner.get_output(task.job_name) + definition.on_failure(task, status, output) + task.status = "failed" + task.updated_at = datetime.now(timezone.utc) + self.db.commit() + runner.delete(task.job_name) + except Exception as error: + log.error("Failed to poll task status for %s: %s", task.job_name, error, exc_info=error) + + result: dict[str, Any] = { + "task_id": task.id, + "status": task.status, + "job_name": task.job_name, + "created_at": task.created_at.isoformat() if task.created_at else None, + "updated_at": task.updated_at.isoformat() if task.updated_at else None, + } + if task.error: + result["error"] = task.error + if task.result: + result["result"] = task.result + + self.set_header("Content-Type", "application/json") + self.write(json.dumps(result)) + + +handlers = [ + (r"/api/beakerhub/admin/node-images/(\d+)/import", NodeImageImportHandler), + (r"/api/beakerhub/admin/node-images/(\d+)/import-status", NodeImageImportStatusHandler), +] + +api_handlers = handlers diff --git a/src/beakerhub/services/task/kubernetes_task_runner.py b/src/beakerhub/services/task/kubernetes_task_runner.py new file mode 100644 index 0000000..dc2d296 --- /dev/null +++ b/src/beakerhub/services/task/kubernetes_task_runner.py @@ -0,0 +1,102 @@ +"""Kubernetes Job implementation of the task-runner service.""" + +from typing import Any + +from traitlets import Dict, Instance, Integer, List, Unicode, default + +from beakerhub.runtimes.kubernetes import ( + KubernetesDefinition, + KubernetesProcess, + KubernetesRuntime, +) +from beakerhub.services.task.base import BaseTaskRunnerService +from beakerhub.tasks.base import BaseImageTask, BaseTaskDefinition + + +class KubernetesTaskRunnerService(BaseTaskRunnerService): + """Adapt Kubernetes Job runtime processes to the task-runner contract.""" + + node_image_resources = Dict( + config=True, + help="Resource requests and limits for the node-image task container.", + ) + backoff_limit = Integer( + 0, + config=True, + help="Number of retries before a task Job is marked as failed.", + ) + active_deadline_seconds = Integer( + 300, + config=True, + help="Maximum time in seconds that a task Job can run.", + ) + ttl_seconds_after_finished = Integer( + 600, + config=True, + help="Time in seconds to retain completed task Jobs.", + ) + node_selector = Dict(config=True, help="Node selector for task Pods.") + tolerations: Any = List(config=True, help="Tolerations for task Pods.") + namespace = Unicode( + config=True, + help="Kubernetes namespace for task Jobs.", + ) + runtime = Instance(KubernetesRuntime, allow_none=False) + + @default("runtime") + def _default_runtime(self) -> KubernetesRuntime: + parent_runtime = getattr(self.parent, "runtime", None) + if isinstance(parent_runtime, KubernetesRuntime): + return parent_runtime + return KubernetesRuntime( + parent=self, + namespace=self._trait_values.get("namespace", "beakerhub"), + node_selector=self._trait_values.get("node_selector", {}), + tolerations=self._trait_values.get("tolerations", []), + ) + + @default("namespace") + def _default_namespace(self) -> str: + return self.runtime.namespace + + @default("node_selector") + def _default_node_selector(self) -> dict[str, str]: + return dict(self.runtime.node_selector) + + @default("tolerations") + def _default_tolerations(self) -> list[dict[str, Any]]: + return list(self.runtime.tolerations) + + def submit(self, task: BaseTaskDefinition) -> KubernetesProcess: + """Submit a supported task as a Kubernetes Job.""" + if not isinstance(task, BaseImageTask): + raise ValueError( + "KubernetesTaskRunnerService only supports image tasks, not " + f"{task.task_type!r}" + ) + task_definition = KubernetesDefinition( + runtime=self.runtime, + namespace=self.namespace, + image=task.image, + entrypoint=tuple(task.entrypoint), + command=tuple(task.command), + working_directory=task.working_directory, + environment=task.environment, + resources=dict(task.resources) or self.node_image_resources, + node_selector=dict(self.node_selector), + tolerations=list(self.tolerations), + labels={"beakerhub/task-type": task.task_type}, + backoff_limit=self.backoff_limit, + active_deadline_seconds=self.active_deadline_seconds, + ttl_seconds_after_finished=self.ttl_seconds_after_finished, + name_prefix=f"beaker-task-{task.task_type.replace('_', '-')}", + ) + return KubernetesProcess.start(task_definition, runtime=self.runtime) + + def get_process(self, external_id: str) -> KubernetesProcess: + """Reconstruct a Job process using the task runner's namespace.""" + return KubernetesProcess( + KubernetesDefinition(runtime=self.runtime, namespace=self.namespace), + runtime=self.runtime, + external_id=external_id, + ) diff --git a/src/beakerhub/spawner/__init__.py b/src/beakerhub/spawner/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/beakerhub/spawner/base.py b/src/beakerhub/spawner/base.py deleted file mode 100644 index d722411..0000000 --- a/src/beakerhub/spawner/base.py +++ /dev/null @@ -1,39 +0,0 @@ -from typing import Any -from uuid import uuid4 - -from jupyterhub.spawner import Spawner -from traitlets import Unicode, default - -class BeakerSpawner(Spawner): - - def __init__(self, **kwargs: Any) -> None: - domain = kwargs.pop("domain", None) - super().__init__(**kwargs) - if domain: - self.proxy_spec = f"{self.name}.{domain}/" - - @property - def name(self) -> str: - name = super().name - return name - - @name.setter - def name(self, value: str) -> str: - if self.orm_spawner: - self.orm_spawner.name = value - return value - - @property - def session_id(self) -> str: - if not self.name: - self.name = str(uuid4()) - return self.name - - def start(self): - raise NotImplementedError() - - def stop(self, now=False): - raise NotImplementedError() - - def poll(self): - raise NotImplementedError() diff --git a/src/beakerhub/tasks/__init__.py b/src/beakerhub/tasks/__init__.py new file mode 100644 index 0000000..2d7a2b0 --- /dev/null +++ b/src/beakerhub/tasks/__init__.py @@ -0,0 +1 @@ +"""Domain task definitions and orchestration.""" diff --git a/src/beakerhub/tasks/base.py b/src/beakerhub/tasks/base.py new file mode 100644 index 0000000..fbd3302 --- /dev/null +++ b/src/beakerhub/tasks/base.py @@ -0,0 +1,54 @@ +"""Runtime-neutral definitions for BeakerHub background tasks.""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Mapping + +from beakerhub.services.task.base import TaskOutput, TaskStatus + +if TYPE_CHECKING: + from beakerhub.orm import BeakerTask + + +class BaseTaskDefinition(ABC): + """A logical task and its task-specific completion behavior. + + Definitions do not contain a runner's external identifier or lifecycle + state. Those belong to the persisted task record and task runner. + """ + + @property + @abstractmethod + def task_type(self) -> str: + """Return the stable type identifier for this task definition.""" + + def on_before_start(self, task: "BeakerTask") -> None: + """Run task-specific preparation after persistence and before submission.""" + + def on_success( + self, + task: "BeakerTask", + status: TaskStatus, + output: TaskOutput, + ) -> None: + """Process captured output after successful completion.""" + + def on_failure( + self, + task: "BeakerTask", + status: TaskStatus, + output: TaskOutput | None, + ) -> None: + """Process diagnostics after unsuccessful completion.""" + + +@dataclass(frozen=True, kw_only=True) +class BaseImageTask(BaseTaskDefinition): + """A task definition that executes a command in an OCI image.""" + + image: str + entrypoint: tuple[str, ...] = () + command: tuple[str, ...] = () + working_directory: str | None = None + environment: Mapping[str, str] = field(default_factory=dict) + resources: Mapping[str, str] = field(default_factory=dict) diff --git a/src/beakerhub/tasks/image_import/__init__.py b/src/beakerhub/tasks/image_import/__init__.py new file mode 100644 index 0000000..139bb3d --- /dev/null +++ b/src/beakerhub/tasks/image_import/__init__.py @@ -0,0 +1,5 @@ +"""Node-image import task.""" + +from .task import launch_import_task + +__all__ = ["launch_import_task"] diff --git a/src/beakerhub/tasks/image_import/task.py b/src/beakerhub/tasks/image_import/task.py new file mode 100644 index 0000000..7e19262 --- /dev/null +++ b/src/beakerhub/tasks/image_import/task.py @@ -0,0 +1,141 @@ +"""Orchestration for importing context metadata from a node image.""" + +import json +import logging +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +from sqlalchemy.orm import Session, object_session + +from beakerhub.orm import NodeImages, NodeImageTask +from beakerhub.services.task.base import TaskOutput, TaskStatus +from beakerhub.tasks.base import BaseImageTask +from beakerhub.utils import ingest_interchange_dump + +if TYPE_CHECKING: + from beakerhub.app import BeakerHub + + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True, kw_only=True) +class ImageImportTask(BaseImageTask): + """Definition and completion behavior for a node-image context import.""" + + node_image: NodeImages + + @classmethod + def from_node_image(cls, node_image: NodeImages) -> "ImageImportTask": + """Create the standard context-dump workload for ``node_image``.""" + return cls( + image=node_image.default_img_string, + entrypoint=("sh", "-c"), + command=("beaker context dump",), + node_image=node_image, + ) + + @property + def task_type(self) -> str: + return "context_import" + + def on_success( + self, + task: NodeImageTask, + status: TaskStatus, + output: TaskOutput, + ) -> None: + """Ingest the context dump into the image's metadata records.""" + db = object_session(task) + if db is None: + raise RuntimeError("Image import task is not attached to a database session") + task.result = ingest_import_output(db, self.node_image, output) + + def on_failure( + self, + task: NodeImageTask, + status: TaskStatus, + output: TaskOutput | None, + ) -> None: + """Store normalized runtime diagnostics on the persisted task.""" + messages = [status.message] + if output and output.stderr: + messages.append(output.stderr) + task.error = ". ".join(message for message in messages if message) + + +def launch_import_task( + db: Session, + app: "BeakerHub", + node_image: NodeImages, +) -> NodeImageTask: + """Create and submit a context-import task for a node image.""" + existing_task = ( + db.query(NodeImageTask) + .filter( + NodeImageTask.node_image_id == node_image.id, + NodeImageTask.task_type == "context_import", + NodeImageTask.status.in_(["pending", "running"]), + ) + .first() + ) + if existing_task: + raise ValueError( + f"Import already in progress for image '{node_image.slug}' " + f"(job: {existing_task.job_name})" + ) + + task = NodeImageTask( + node_image_id=node_image.id, + task_type="context_import", + status="pending", + ) + db.add(task) + db.commit() + + try: + running_task = app.task_runner.submit(ImageImportTask.from_node_image(node_image)) + except Exception as error: + task.status = "failed" + task.error = f"Failed to create task: {error}" + task.updated_at = datetime.now(timezone.utc) + db.commit() + log.exception("Failed to create import task for %s", node_image.slug) + raise RuntimeError(f"Failed to create import task: {error}") from error + + task.job_name = running_task.external_id + task.status = "running" + task.updated_at = datetime.now(timezone.utc) + db.commit() + log.info( + "Launched import task %s for node image %s", + running_task.external_id, + node_image.slug, + ) + return task + + +def ingest_import_output( + db: Session, + node_image: NodeImages, + output: TaskOutput, +) -> dict[str, int]: + """Ingest the context dumps written to an image-import task's stdout.""" + + body = json.loads(output.stdout) + if not isinstance(body, list): + body = [body] + + combined_stats: dict[str, int] = {} + for dump in body: + stats = ingest_interchange_dump( + db=db, + dump=dump, + node_image=node_image, + enable_contexts=True, + preserve_curated=True, + ) + for key, value in stats.items(): + combined_stats[key] = combined_stats.get(key, 0) + value + return combined_stats diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 03910a5..9c2daad 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -3,11 +3,26 @@ # SPDX-License-Identifier: MIT """Shared fixtures for unit tests.""" -import pytest +from unittest.mock import Mock + import boto3 +import pytest +from kubernetes import config as k8s_config +from kubernetes.config.config_exception import ConfigException from moto import mock_aws +@pytest.fixture(autouse=True) +def mock_kubernetes_config(monkeypatch): + """Prevent unit tests from requiring an in-cluster config or kubeconfig.""" + monkeypatch.setattr( + k8s_config, + "load_incluster_config", + Mock(side_effect=ConfigException("Not running in a cluster")), + ) + monkeypatch.setattr(k8s_config, "load_kube_config", Mock()) + + @pytest.fixture def cognito_user_pool(): """ diff --git a/src/beakerhub/nodes/__init__.py b/tests/unit/runtimes/__init__.py similarity index 100% rename from src/beakerhub/nodes/__init__.py rename to tests/unit/runtimes/__init__.py diff --git a/tests/unit/runtimes/test_aws_ecs.py b/tests/unit/runtimes/test_aws_ecs.py new file mode 100644 index 0000000..ec4b54e --- /dev/null +++ b/tests/unit/runtimes/test_aws_ecs.py @@ -0,0 +1,287 @@ +"""Tests for ECS runtime process construction and control-plane behavior.""" + +import re +from unittest.mock import Mock + +import boto3 +import pytest + +from beakerhub.runtimes.aws_ecs import ( + AwsEcsDefinition, + AwsEcsProcess, + AwsEcsRuntime, +) + + +def ecs_runtime(**overrides): + ecs = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + logs = boto3.client( + "logs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + values = { + "ecs_client": ecs, + "logs_client": logs, + "cluster_name": "runtime-cluster", + "subnets": ["subnet-runtime"], + "security_groups": ["sg-runtime"], + "log_group": "/beakerhub/runtime", + "launch_options": {"enableExecuteCommand": True}, + } + values.update(overrides) + return AwsEcsRuntime(**values) + + +def test_definition_uses_runtime_defaults_and_allows_definition_overrides(): + runtime = ecs_runtime() + definition = AwsEcsDefinition( + image="example:latest", + task_definition_arn="notebook:1", + runtime=runtime, + ) + overridden = AwsEcsDefinition( + image="example:latest", + task_definition_arn="notebook:1", + runtime=runtime, + cluster_name="definition-cluster", + subnets=["subnet-definition"], + ) + + assert definition.cluster_name == "runtime-cluster" + assert definition.subnets == ["subnet-runtime"] + assert definition.security_groups == ["sg-runtime"] + assert definition.log_group == "/beakerhub/runtime" + assert definition.launch_options == {"enableExecuteCommand": True} + assert overridden.cluster_name == "definition-cluster" + assert overridden.subnets == ["subnet-definition"] + + definition.subnets.append("local-change") + assert runtime.subnets == ["subnet-runtime"] + + +def test_dynamic_task_definition_family_is_ecs_safe_and_unique(): + runtime = ecs_runtime(task_definition_name="beaker.hub/tasks") + digest = "a" * 300 + definition = AwsEcsDefinition( + image=f"registry.example/notebook@sha256:{digest}", + runtime=runtime, + ) + similar = AwsEcsDefinition( + image="registry.example/notebook.sha256:latest", + runtime=runtime, + ) + + assert len(definition.task_definition_name) <= 255 + assert re.fullmatch(r"[A-Za-z0-9_-]+", definition.task_definition_name) + assert definition.task_definition_name != similar.task_definition_name + + +def test_fixed_definition_launch_uses_runtime_networking_and_process_metadata(): + runtime = ecs_runtime() + runtime.ecs_client.run_task = Mock( + return_value={"tasks": [{"taskArn": "task-arn"}]} + ) + runtime.ecs_client.register_task_definition = Mock() + + process = AwsEcsProcess.start( + AwsEcsDefinition( + image="notebook:latest", + task_definition_arn="notebook:1", + container_name="notebook", + environment={"BEAKERHUB_USER": "ada"}, + tags={"beaker-session": "session-1"}, + ), + process_type="service", + runtime=runtime, + ) + + assert process.external_id == "task-arn" + runtime.ecs_client.register_task_definition.assert_not_called() + request = runtime.ecs_client.run_task.call_args.kwargs + assert request["cluster"] == "runtime-cluster" + assert request["networkConfiguration"] == { + "awsvpcConfiguration": { + "subnets": ["subnet-runtime"], + "securityGroups": ["sg-runtime"], + "assignPublicIp": "DISABLED", + } + } + assert request["overrides"]["containerOverrides"] == [ + { + "name": "notebook", + "environment": [{"name": "BEAKERHUB_USER", "value": "ada"}], + } + ] + assert {tag["key"]: tag["value"] for tag in request["tags"]} == { + "beaker-session": "session-1", + "beakerhub-process": "true", + "beakerhub-process-type": "service", + } + + +def test_dynamic_definition_includes_volumes_and_container_mount_points(): + definition = AwsEcsDefinition( + image="example:latest", + runtime=ecs_runtime(), + volumes=[ + { + "name": "user-storage", + "efsVolumeConfiguration": {"fileSystemId": "fs-12345678"}, + } + ], + mount_points=[ + { + "sourceVolume": "user-storage", + "containerPath": "/home/beaker", + "readOnly": False, + } + ], + ) + + task_definition = definition.aws_task_definition + + assert task_definition["volumes"] == definition.volumes + assert task_definition["containerDefinitions"][0]["mountPoints"] == ( + definition.mount_points + ) + + +def test_dynamic_definition_resources_override_runtime_defaults(): + runtime = ecs_runtime( + default_task_cpu="2 vcpu", + default_task_memory="8GB", + default_task_ephemeral_storage=40, + ) + runtime.ecs_client.list_task_definitions = Mock( + return_value={"taskDefinitionArns": []} + ) + runtime.ecs_client.register_task_definition = Mock( + return_value={"taskDefinition": {"taskDefinitionArn": "definition-arn"}} + ) + runtime.ecs_client.run_task = Mock( + return_value={"tasks": [{"taskArn": "task-arn"}]} + ) + process = AwsEcsProcess.start( + AwsEcsDefinition( + image="example:latest", + cpu="1 vcpu", + memory="2GB", + ephemeral_storage=25, + ), + runtime=runtime, + ) + + assert process.external_id == "task-arn" + + request = runtime.ecs_client.register_task_definition.call_args.kwargs + assert request["cpu"] == "1 vcpu" + assert request["memory"] == "2GB" + assert request["ephemeralStorage"] == {"sizeInGiB": 25} + + +def test_definition_reuses_a_matching_task_definition(): + runtime = ecs_runtime() + definition = AwsEcsDefinition(image="example:latest", runtime=runtime) + existing = { + **definition.aws_task_definition, + "taskDefinitionArn": "definition-arn", + "revision": 1, + "status": "ACTIVE", + "compatibilities": ["EC2", "FARGATE"], + "requiresAttributes": [], + } + runtime.ecs_client.list_task_definitions = Mock( + return_value={"taskDefinitionArns": ["definition-arn"]} + ) + runtime.ecs_client.describe_task_definition = Mock( + return_value={"taskDefinition": existing} + ) + runtime.ecs_client.register_task_definition = Mock() + + assert definition.find_or_register_task_definition() == "definition-arn" + runtime.ecs_client.register_task_definition.assert_not_called() + + + +@pytest.mark.parametrize( + ("runtime_values", "message"), + [ + ({"cluster_name": ""}, "cluster_name"), + ({"subnets": []}, "subnets"), + ({"launch_type": "EC2"}, "ephemeral_storage"), + ( + { + "launch_type": "MANAGED_INSTANCES", + "default_task_ephemeral_storage": 0, + }, + "capacityProviderStrategy", + ), + ], +) +def test_definition_validates_its_runtime_derived_configuration( + runtime_values, + message, +): + definition = AwsEcsDefinition( + image="example:latest", + runtime=ecs_runtime(**runtime_values), + ) + + with pytest.raises(ValueError, match=message): + definition.validate_configuration() + + +@pytest.mark.parametrize( + ("exit_code", "state"), + [(0, "completed"), (23, "failed")], +) +def test_process_status_preserves_container_exit_code(exit_code, state): + runtime = ecs_runtime() + runtime.ecs_client.describe_tasks = Mock( + return_value={ + "tasks": [ + { + "lastStatus": "STOPPED", + "containers": [ + {"name": "task-container", "exitCode": exit_code} + ], + } + ] + } + ) + process = AwsEcsProcess( + AwsEcsDefinition(image="example:latest", runtime=runtime), + runtime=runtime, + external_id="task-arn", + ) + + status = process.describe() + + assert status.state == state + assert status.exit_code == exit_code + + +def test_runtime_reads_cloudwatch_log_pages(): + runtime = ecs_runtime() + runtime.logs_client.get_log_events = Mock( + side_effect=[ + { + "events": [{"message": "first"}], + "nextForwardToken": "next", + }, + { + "events": [{"message": "second"}], + "nextForwardToken": "next", + }, + ] + ) + + assert runtime.get_log_events("/logs", "stream") == ["first", "second"] + assert runtime.logs_client.get_log_events.call_count == 2 diff --git a/tests/unit/runtimes/test_base.py b/tests/unit/runtimes/test_base.py new file mode 100644 index 0000000..73940b6 --- /dev/null +++ b/tests/unit/runtimes/test_base.py @@ -0,0 +1,73 @@ +"""Tests for provider-neutral runtime contracts.""" + +from traitlets import Unicode + +from beakerhub.runtimes.base import ( + BaseDefinition, + BaseProcess, + BaseRuntime, + BaseRuntimeBundle, + ProcessStatus, +) + + +class ExampleDefinition(BaseDefinition): + value = Unicode("example", config=True) + + +class ExampleRuntime(BaseRuntime): + pass + + +class ExampleProcess(BaseProcess): + @classmethod + def start(cls, definition, **kwargs): + return cls(definition, external_id="example-process", **kwargs) + + def describe(self): + return ProcessStatus("running") + + def collect_output(self): + return None + + def stop(self): + return None + + + +class ExampleBundle(BaseRuntimeBundle): + runtime_class = ExampleRuntime + process_class = ExampleProcess + definition_class = ExampleDefinition + + # default_dashboard_class = + # default_spawner_class = + # default_task_runner_class = + +def test_bundle_creates_definition_and_process_with_application_runtime(): + bundle = ExampleBundle() + runtime = ExampleRuntime() + + definition = bundle.create_definition(runtime=runtime, value="configured") + process = bundle.start_process(definition, runtime=runtime) + + assert isinstance(definition, ExampleDefinition) + assert definition.value == "configured" + assert definition.parent is runtime + assert definition.runtime is runtime + assert isinstance(process.runtime, ExampleRuntime) + assert process.parent is runtime + assert process.runtime is runtime + assert process.external_id == "example-process" + assert process.status == ProcessStatus("running") + + +def test_process_status_marks_only_terminal_states_done(): + assert ProcessStatus("pending").done is False + assert ProcessStatus("running").done is False + assert ProcessStatus("completed").done is True + assert ProcessStatus("failed").done is True + + +def test_process_status_retains_an_optional_exit_code(): + assert ProcessStatus("failed", exit_code=23).exit_code == 23 diff --git a/tests/unit/runtimes/test_kubernetes.py b/tests/unit/runtimes/test_kubernetes.py new file mode 100644 index 0000000..a683f28 --- /dev/null +++ b/tests/unit/runtimes/test_kubernetes.py @@ -0,0 +1,113 @@ +"""Tests for Kubernetes runtime Job process construction.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from beakerhub.runtimes.kubernetes import ( + KubernetesDefinition, + KubernetesProcess, + KubernetesRuntime, +) + + +def runtime(**overrides): + values = { + "namespace": "runtime-namespace", + "node_selector": {"pool": "compute"}, + "tolerations": [{"key": "dedicated", "operator": "Exists"}], + "service_account": "runtime-workload", + "base_labels": {"environment": "test"}, + } + values.update(overrides) + return KubernetesRuntime(**values) + + +def test_job_inherits_runtime_pod_defaults_and_base_labels(): + process = KubernetesProcess( + KubernetesDefinition(image="example:latest"), + runtime=runtime(), + external_id="beaker-process-123", + ) + + job = process._job() + pod_spec = job.spec.template.spec + + assert job.metadata.namespace == "runtime-namespace" + assert job.metadata.labels["environment"] == "test" + assert job.spec.template.metadata.labels["environment"] == "test" + assert pod_spec.node_selector == {"pool": "compute"} + assert pod_spec.tolerations[0].key == "dedicated" + assert pod_spec.service_account_name == "runtime-workload" + + +def test_definition_can_override_or_disable_runtime_pod_defaults(): + process = KubernetesProcess( + KubernetesDefinition( + image="example:latest", + namespace="definition-namespace", + node_selector={}, + tolerations=(), + service_account="", + labels={"environment": "definition"}, + ), + runtime=runtime(), + external_id="beaker-process-123", + ) + + job = process._job() + pod_spec = job.spec.template.spec + + assert job.metadata.namespace == "definition-namespace" + assert job.metadata.labels["environment"] == "definition" + assert pod_spec.node_selector is None + assert pod_spec.tolerations is None + assert pod_spec.service_account_name is None + + +def test_start_creates_job_and_records_the_provider_identifier(monkeypatch): + batch_api = Mock() + core_api = Mock() + instance = runtime() + instance.batch_api = batch_api + instance.core_api = core_api + + process = KubernetesProcess.start( + KubernetesDefinition(image="example:latest", name_prefix="image-import"), + runtime=instance, + ) + + assert process.external_id.startswith("image-import-") + request = batch_api.create_namespaced_job.call_args.kwargs + assert request["namespace"] == "runtime-namespace" + assert request["body"].metadata.name == process.external_id + + +def test_service_processes_are_not_mistaken_for_jobs(): + with pytest.raises(NotImplementedError, match="run-to-completion Jobs"): + KubernetesProcess.start( + KubernetesDefinition(image="example:latest"), + runtime=runtime(), + process_type="service", + ) + + +def test_describe_reports_job_completion(monkeypatch): + batch_api = Mock() + batch_api.read_namespaced_job.return_value = SimpleNamespace( + status=SimpleNamespace(succeeded=1, failed=None, active=None) + ) + instance = runtime() + instance.batch_api = batch_api + instance.core_api = Mock() + process = KubernetesProcess( + KubernetesDefinition(image="example:latest"), + runtime=instance, + external_id="beaker-process-123", + ) + + status = process.describe() + + assert status.state == "completed" + assert status.message == "Job completed successfully" diff --git a/tests/unit/services/spawner/__init__.py b/tests/unit/services/spawner/__init__.py new file mode 100644 index 0000000..d024455 --- /dev/null +++ b/tests/unit/services/spawner/__init__.py @@ -0,0 +1 @@ +"""Tests for spawner services.""" diff --git a/tests/unit/services/spawner/test_aws_ecs.py b/tests/unit/services/spawner/test_aws_ecs.py new file mode 100644 index 0000000..2677aa2 --- /dev/null +++ b/tests/unit/services/spawner/test_aws_ecs.py @@ -0,0 +1,563 @@ +"""Tests for the ECS runtime-backed JupyterHub spawner.""" + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import boto3 +import pytest + +from beakerhub.runtimes.aws_ecs import AwsEcsRuntime +from beakerhub.runtimes.base import ProcessStatus +from beakerhub.services.spawner.aws_ecs_spawner import BeakerAwsECSSpawner + + +TASK_ARN = "arn:aws:ecs:us-east-1:123456789012:task/test/task-123" + + +def runtime(**overrides) -> AwsEcsRuntime: + ecs = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + values = { + "ecs_client": ecs, + "cluster_name": "runtime-cluster", + "subnets": ["subnet-runtime"], + "security_groups": ["sg-runtime"], + "launch_options": {"enableExecuteCommand": True}, + "task_overrides": {"cpu": "512"}, + "log_group": "/beakerhub/runtime", + "log_stream_prefix": "runtime-session", + "execution_role_arn": "arn:aws:iam::123456789012:role/execution", + "task_role_arn": "arn:aws:iam::123456789012:role/task", + "task_group": "runtime-group", + "default_task_cpu": "2 vcpu", + "default_task_memory": "8GB", + "default_task_ephemeral_storage": 40, + } + values.update(overrides) + return AwsEcsRuntime(**values) + + +def configured_spawner(**overrides): + """A minimal object sufficient to exercise unbound spawner methods.""" + shared_runtime = overrides.get("runtime", runtime()) + values = { + "runtime": shared_runtime, + "user": SimpleNamespace(id=1, name="ada"), + "cluster_name": "notebooks", + "task_definition": "beaker-notebook:3", + "container_name": "notebook", + "container_tags": {"environment": "test"}, + "session_id": "session-123", + "subnets": ["subnet-123"], + "security_groups": ["sg-123"], + "assign_public_ip": True, + "launch_options": {"enableExecuteCommand": True}, + "efs_volume_configuration": {}, + "efs_mount_path": "", + "efs_volume_name": "beakerhub-efs", + "task_overrides": dict(shared_runtime.task_overrides or {}), + "log_group": shared_runtime.log_group, + "log_stream_prefix": shared_runtime.log_stream_prefix, + "execution_role_arn": shared_runtime.execution_role_arn, + "task_role_arn": shared_runtime.task_role_arn, + "cpu_architecture": shared_runtime.cpu_architecture, + "task_group": shared_runtime.task_group, + "launch_type": shared_runtime.launch_type, + "default_task_cpu": shared_runtime.default_task_cpu, + "default_task_memory": shared_runtime.default_task_memory, + "default_task_ephemeral_storage": shared_runtime.default_task_ephemeral_storage, + "volume_configs": [], + "image": "registry.example/beaker-node:latest", + "get_env": lambda: {"BEAKERHUB_USER": "ada", "NUMBER": 2}, + "_notebook_config": lambda: "", + "task_arn": None, + "debug": False, + "log": Mock(), + } + values.update(overrides) + values.setdefault( + "task_definition_name", + "beakerhub-notebook_d25e4b5db3dca83c3aae6b9fe0d9c3e9fc019ced629567804f94092004ac0316", + ) + spawner = SimpleNamespace(**values) + spawner._task_tags = lambda: BeakerAwsECSSpawner._task_tags(spawner) + spawner._efs_volumes = lambda: BeakerAwsECSSpawner._efs_volumes(spawner) + spawner._efs_access_point_id = ( + lambda file_system_id, user_id, user_subdir: + BeakerAwsECSSpawner._efs_access_point_id( + spawner, file_system_id, user_id, user_subdir + ) + ) + spawner._find_efs_access_point = BeakerAwsECSSpawner._find_efs_access_point + spawner._efs_access_point_token = BeakerAwsECSSpawner._efs_access_point_token + spawner._wait_for_efs_access_point = BeakerAwsECSSpawner._wait_for_efs_access_point + spawner._task_overrides = lambda: BeakerAwsECSSpawner._task_overrides(spawner) + spawner._config_writer_container = ( + lambda: BeakerAwsECSSpawner._config_writer_container(spawner) + ) + spawner._definition = lambda: BeakerAwsECSSpawner._definition(spawner) + spawner._process = lambda: BeakerAwsECSSpawner._process(spawner) + spawner._wait_for_task_ip = AsyncMock(return_value="10.0.0.42") + return spawner + + +def test_definition_contains_session_overrides_and_runtime_launch_options(): + spawner = configured_spawner( + volume_configs=[{"name": "session-volume", "managedEBSVolume": {}}] + ) + + definition = BeakerAwsECSSpawner._definition(spawner) + + assert definition.runtime is spawner.runtime + assert definition.image == "registry.example/beaker-node:latest" + assert definition.task_definition_arn == "beaker-notebook:3" + assert definition.container_name == "notebook" + assert definition.environment == {} + assert definition.task_overrides == { + "cpu": "512", + "containerOverrides": [ + { + "name": "notebook", + "environment": [ + {"name": "BEAKERHUB_USER", "value": "ada"}, + {"name": "NUMBER", "value": "2"}, + ], + } + ] + } + assert definition.tags == { + "environment": "test", + "beaker-session": "session-123", + } + assert definition.cluster_name == "notebooks" + assert definition.subnets == ["subnet-123"] + assert definition.security_groups == ["sg-123"] + assert definition.assign_public_ip is True + assert definition.log_group == "/beakerhub/runtime" + assert definition.log_stream_prefix == "runtime-session" + assert definition.execution_role_arn.endswith("role/execution") + assert definition.task_role_arn.endswith("role/task") + assert definition.cpu_architecture == "X86_64" + assert definition.task_group == "runtime-group" + assert definition.launch_type == "FARGATE" + assert definition.cpu == "2 vcpu" + assert definition.memory == "8GB" + assert definition.ephemeral_storage == 40 + assert definition.launch_options == { + "enableExecuteCommand": False, + "volumeConfigurations": [{"name": "session-volume", "managedEBSVolume": {}}], + } + + +def test_definition_enables_ecs_exec_when_spawner_debug_is_enabled(): + spawner = configured_spawner(debug=True) + + definition = BeakerAwsECSSpawner._definition(spawner) + + assert definition.launch_options["enableExecuteCommand"] is True + + +def test_definition_adds_an_optional_efs_volume_and_mount(monkeypatch): + shared_runtime = runtime() + tagging = Mock() + tagging.get_resources.return_value = {"ResourceTagMappingList": []} + efs = Mock() + efs.create_access_point.return_value = {"AccessPointId": "fsap-12345678"} + efs.describe_access_points.return_value = { + "AccessPoints": [{"LifeCycleState": "available"}] + } + boto3_client = boto3.client + + def aws_client(service_name, **kwargs): + if service_name == "efs": + return efs + if service_name == "resourcegroupstaggingapi": + return tagging + return boto3_client(service_name, **kwargs) + + monkeypatch.setattr( + "beakerhub.services.spawner.aws_ecs_spawner.boto3.client", + aws_client, + ) + spawner = configured_spawner( + runtime=shared_runtime, + efs_volume_configuration={ + "fileSystemId": "fs-12345678", + "rootDirectory": "/notebooks", + "transitEncryption": "ENABLED", + }, + efs_mount_path="/home/beaker", + efs_volume_name="user-storage", + ) + + volumes, mount_points = BeakerAwsECSSpawner._efs_volumes(spawner) + + assert volumes == [ + { + "name": "user-storage", + "efsVolumeConfiguration": { + "fileSystemId": "fs-12345678", + "transitEncryption": "ENABLED", + "authorizationConfig": {"accessPointId": "fsap-12345678"}, + }, + } + ] + assert mount_points == [ + { + "sourceVolume": "user-storage", + "containerPath": "/home/beaker", + "readOnly": False, + } + ] + efs.create_access_point.assert_called_once_with( + ClientToken="580d2dbbb845c84be8c83a0f6188bbf573b7bed4850770030bb762ed79e89203", + FileSystemId="fs-12345678", + RootDirectory={ + "Path": "/user-storage/ada---fdee430d", + "CreationInfo": { + "OwnerUid": 1000, + "OwnerGid": 1000, + "Permissions": "755", + }, + }, + Tags=[ + {"Key": "beakerhub-user-id", "Value": "1"}, + {"Key": "beakerhub-user-name", "Value": "ada"}, + { + "Key": "beakerhub-user-subdir", + "Value": "/user-storage/ada---fdee430d", + }, + ], + ) + + +def test_dynamic_definition_includes_config_writer_sidecar(): + spawner = configured_spawner( + task_definition="", + _notebook_config=lambda: "c.Foo.bar = 'baz'", + ) + + definition = BeakerAwsECSSpawner._definition(spawner) + + assert definition.sidecar_containers == [ + { + "name": "config-writer", + "image": "busybox:latest", + "command": [ + "/bin/sh", + "-ec", + "printf '%s' \"$BEAKER_NOTEBOOK_CONFIG\" > " + "/opt/beaker/config/beaker_config.py", + ], + "environment": [ + {"name": "BEAKER_NOTEBOOK_CONFIG", "value": "c.Foo.bar = 'baz'"} + ], + "mountPoints": [ + { + "containerPath": "/opt/beaker/config", + "sourceVolume": "config", + } + ], + "essential": False, + } + ] + assert definition.aws_task_definition["containerDefinitions"][0]["dependsOn"] == [ + {"containerName": "config-writer", "condition": "SUCCESS"} + ] + assert definition.aws_task_definition["containerDefinitions"][1] == ( + definition.sidecar_containers[0] + ) + + +def test_session_tag_cannot_be_overridden_by_configuration(): + spawner = configured_spawner(container_tags={"beaker-session": "incorrect"}) + + assert BeakerAwsECSSpawner._task_tags(spawner) == { + "beaker-session": "session-123" + } + + +@pytest.mark.parametrize( + "overrides", + [ + {"efs_volume_configuration": {"fileSystemId": "fs-123"}}, + {"efs_mount_path": "/home/beaker"}, + { + "efs_volume_configuration": {"rootDirectory": "/"}, + "efs_mount_path": "/home/beaker", + }, + ], +) +def test_definition_rejects_incomplete_efs_configuration(overrides): + spawner = configured_spawner(**overrides) + + with pytest.raises(ValueError, match="efs_"): + BeakerAwsECSSpawner._definition(spawner) + + +def test_definition_rejects_an_empty_container_name(): + spawner = configured_spawner(container_name="") + + with pytest.raises(ValueError, match="container_name"): + BeakerAwsECSSpawner._definition(spawner) + + +async def test_start_uses_the_runtime_process_and_persists_its_arn(): + shared_runtime = runtime() + shared_runtime.ecs_client.run_task = Mock( + return_value={"tasks": [{"taskArn": TASK_ARN}]} + ) + shared_runtime.ecs_client.register_task_definition = Mock() + spawner = configured_spawner(runtime=shared_runtime) + + await BeakerAwsECSSpawner.start(spawner) + + assert spawner.task_arn == TASK_ARN + request = shared_runtime.ecs_client.run_task.call_args.kwargs + assert request["taskDefinition"] == "beaker-notebook:3" + assert request["cluster"] == "notebooks" + assert request["networkConfiguration"] == { + "awsvpcConfiguration": { + "subnets": ["subnet-123"], + "securityGroups": ["sg-123"], + "assignPublicIp": "ENABLED", + } + } + assert {tag["key"]: tag["value"] for tag in request["tags"]} == { + "environment": "test", + "beaker-session": "session-123", + "beakerhub-process": "true", + "beakerhub-process-type": "service", + } + shared_runtime.ecs_client.register_task_definition.assert_not_called() + + +async def test_start_registers_and_reuses_an_image_specific_definition(): + shared_runtime = runtime() + shared_runtime.ecs_client.list_task_definitions = Mock( + return_value={"taskDefinitionArns": []} + ) + shared_runtime.ecs_client.register_task_definition = Mock( + return_value={"taskDefinition": {"taskDefinitionArn": "definition-arn"}} + ) + shared_runtime.ecs_client.run_task = Mock( + return_value={"tasks": [{"taskArn": TASK_ARN}]} + ) + spawner = configured_spawner(runtime=shared_runtime, task_definition="") + + await BeakerAwsECSSpawner.start(spawner) + + registered = shared_runtime.ecs_client.register_task_definition.call_args.kwargs + assert registered["family"] == ( + "beakerhub-notebook_" + "d25e4b5db3dca83c3aae6b9fe0d9c3e9fc019ced629567804f94092004ac0316" + ) + assert registered["containerDefinitions"][0]["image"] == ( + "registry.example/beaker-node:latest" + ) + assert "environment" not in registered["containerDefinitions"][0] + existing = { + **registered, + "taskDefinitionArn": "definition-arn", + "revision": 1, + "status": "ACTIVE", + } + shared_runtime.ecs_client.list_task_definitions = Mock( + return_value={"taskDefinitionArns": ["definition-arn"]} + ) + shared_runtime.ecs_client.describe_task_definition = Mock( + return_value={"taskDefinition": existing} + ) + shared_runtime.ecs_client.register_task_definition.reset_mock() + + await BeakerAwsECSSpawner.start(spawner) + + shared_runtime.ecs_client.register_task_definition.assert_not_called() + + +async def test_start_validates_fargate_networking_from_the_shared_runtime(): + shared_runtime = runtime(subnets=[]) + spawner = configured_spawner(runtime=shared_runtime, subnets=[]) + + with pytest.raises(ValueError, match="subnets"): + await BeakerAwsECSSpawner.start(spawner) + + +async def test_capacity_provider_launch_options_omit_launch_type(): + shared_runtime = runtime() + shared_runtime.ecs_client.run_task = Mock( + return_value={"tasks": [{"taskArn": TASK_ARN}]} + ) + spawner = configured_spawner( + runtime=shared_runtime, + launch_options={ + "capacityProviderStrategy": [ + {"capacityProvider": "FARGATE", "weight": 1} + ] + }, + ) + + await BeakerAwsECSSpawner.start(spawner) + + request = shared_runtime.ecs_client.run_task.call_args.kwargs + assert "launchType" not in request + assert request["capacityProviderStrategy"] == [ + {"capacityProvider": "FARGATE", "weight": 1} + ] + + +async def test_wait_for_task_ip_raises_when_the_task_stops(): + shared_runtime = runtime() + shared_runtime.ecs_client.describe_tasks = Mock( + return_value={ + "tasks": [ + { + "lastStatus": "STOPPED", + "stoppedReason": "CannotPullContainerError", + "containers": [ + {"name": "notebook", "reason": "Image pull failed"} + ], + } + ] + } + ) + spawner = configured_spawner(runtime=shared_runtime) + process = SimpleNamespace( + definition=SimpleNamespace(cluster_name="notebooks", container_name="notebook"), + external_id=TASK_ARN, + ) + + with pytest.raises(RuntimeError, match="Image pull failed"): + await BeakerAwsECSSpawner._wait_for_task_ip(spawner, process) + + +async def test_stop_waits_for_ecs_to_deprovision_the_session_task(): + process = SimpleNamespace( + stop=Mock(), + await_completion=AsyncMock(), + ) + spawner = configured_spawner(task_arn=TASK_ARN) + spawner._process = lambda: process + + await BeakerAwsECSSpawner.stop(spawner) + + process.stop.assert_called_once_with() + process.await_completion.assert_awaited_once_with() + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (ProcessStatus("pending"), None), + (ProcessStatus("running"), None), + (ProcessStatus("completed", exit_code=0), 0), + (ProcessStatus("failed", exit_code=23), 23), + (ProcessStatus("completed"), 0), + (ProcessStatus("failed"), 1), + ], +) +async def test_poll_maps_runtime_status_to_jupyterhub_semantics(status, expected): + spawner = configured_spawner(task_arn=TASK_ARN) + spawner._process = lambda: SimpleNamespace(status=status) + + assert await BeakerAwsECSSpawner.poll(spawner) == expected + + +async def test_poll_reports_stopped_when_no_session_task_was_started(): + spawner = configured_spawner(task_arn=None) + + assert await BeakerAwsECSSpawner.poll(spawner) == 0 + + +def test_state_round_trip_restores_the_session_task_arn(): + shared_runtime = runtime() + user = SimpleNamespace(settings={"app": SimpleNamespace(runtime=shared_runtime)}) + source = BeakerAwsECSSpawner( + user=user, + task_definition="beaker-notebook:3", + container_name="notebook", + ) + source.task_arn = TASK_ARN + restored = BeakerAwsECSSpawner( + user=user, + task_definition="beaker-notebook:3", + container_name="notebook", + ) + + restored.load_state(source.get_state()) + + assert restored.runtime is shared_runtime + assert restored.task_arn == TASK_ARN + + +def test_notebook_config_configures_the_notebook_service(monkeypatch): + from beakerhub.app import BeakerHub + from traitlets.config import Config + + class NotebookConfigSpawner: + log = Mock() + + monkeypatch.setattr( + BeakerHub, + "instance", + classmethod(lambda cls: SimpleNamespace(subdomain_host="notebooks.example.test")), + ) + BeakerAwsECSSpawner._notebook_config.cache_clear() + try: + config_text = BeakerAwsECSSpawner._notebook_config(NotebookConfigSpawner()) + finally: + BeakerAwsECSSpawner._notebook_config.cache_clear() + + config = Config() + exec(config_text, {"get_config": lambda: config}) + + assert config.BaseBeakerApp.allow_origin_pat == "^notebooks.example.test$" + assert config.BaseBeakerApp.tornado_settings == { + "headers": { + "Content-Security-Policy": ( + "frame-ancestors 'self' notebooks.example.test" + ) + } + } + assert config.BaseBeakerApp.identity_provider_class == ( + "beakerhub.auth.node.BeakerhubNodeIdentityProvider" + ) + assert config.BaseBeakerApp.authorizer_class == ( + "beakerhub.auth.node.BeakerhubNodeAuthorizer" + ) + assert config.BaseBeakerApp.secrets_manager_class == ( + "beakerhub.services.secrets.beakerhub.BeakerhubSecretsManager" + ) + assert config.BeakerhubSecretsManager.policy_override_env_key_suffix == "_secret_policy" + assert config.FileNotebookManager.notebook_path == ".notebooks" + assert config.FileNotebookManager.snapshot_path == ".notebooks" + + +def test_default_runtime_uses_the_application_runtime(): + shared_runtime = runtime() + user = SimpleNamespace( + id=1, + settings={"app": SimpleNamespace(runtime=shared_runtime)}, + ) + + spawner = BeakerAwsECSSpawner( + user=user, + image="registry.example/beaker-node:latest", + task_definition="beaker-notebook:3", + container_name="notebook", + ) + + assert spawner.runtime is shared_runtime + assert spawner.log_group == "/beakerhub/runtime" + assert spawner.execution_role_arn.endswith("role/execution") + assert spawner.task_group == "runtime-group" + assert spawner.default_task_cpu == "2 vcpu" + assert spawner.task_definition_name == ( + "beakerhub-notebook_" + "d25e4b5db3dca83c3aae6b9fe0d9c3e9fc019ced629567804f94092004ac0316" + ) diff --git a/tests/unit/spawner/test_base.py b/tests/unit/services/spawner/test_base.py similarity index 97% rename from tests/unit/spawner/test_base.py rename to tests/unit/services/spawner/test_base.py index 0beba05..f1dc9b8 100644 --- a/tests/unit/spawner/test_base.py +++ b/tests/unit/services/spawner/test_base.py @@ -1,14 +1,14 @@ # SPDX-FileCopyrightText: 2024-present Jataware Corp # # SPDX-License-Identifier: MIT -"""Unit tests for beakerhub.spawner.base module.""" +"""Unit tests for the base spawner service.""" import re from unittest.mock import MagicMock, patch, PropertyMock import pytest -from beakerhub.spawner.base import BeakerSpawner +from beakerhub.services.spawner.base import BeakerSpawner class TestBeakerSpawner: diff --git a/tests/unit/spawner/test_kubernetes.py b/tests/unit/services/spawner/test_kubernetes.py similarity index 86% rename from tests/unit/spawner/test_kubernetes.py rename to tests/unit/services/spawner/test_kubernetes.py index 0782009..5b36f4a 100644 --- a/tests/unit/spawner/test_kubernetes.py +++ b/tests/unit/services/spawner/test_kubernetes.py @@ -1,13 +1,44 @@ # SPDX-FileCopyrightText: 2024-present Jataware Corp # # SPDX-License-Identifier: MIT -"""Unit tests for beakerhub.spawner.kubernetes module.""" +"""Unit tests for the Kubernetes spawner service.""" from unittest.mock import MagicMock, patch, PropertyMock import pytest -from beakerhub.spawner.kubernetes import BeakerKubeSpawner +from beakerhub.services.spawner.kubernetes_spawner import BeakerKubeSpawner + + +def spawner_for_apply_user_options(**overrides): + """Construct a minimal real spawner for unbound method tests. + + The method now uses ``super()``, so its receiver must be an actual + ``BeakerKubeSpawner`` rather than a mock with that specification. + """ + spawner = object.__new__(BeakerKubeSpawner) + spawner._trait_values = {} + spawner._trait_notifiers = {} + spawner._trait_validators = {} + spawner._cross_validation_lock = False + + db = MagicMock() + db.query.return_value.filter.return_value.first.return_value = None + db.query.return_value.filter.return_value.all.return_value = [] + values = { + "db": db, + "beaker_context": "", + "context_config": {}, + "node_env": {}, + "node_policy_overrides": {}, + "debug": False, + "default_registry": "registry.example", + "default_tag": "latest", + } + values.update(overrides) + for name, value in values.items(): + setattr(spawner, name, value) + return spawner class TestBeakerKubeSpawner: @@ -126,82 +157,75 @@ def test_get_env_preserves_parent_env(self, mock_spawner, mock_user): class TestApplyUserOptions: - """Tests for BeakerKubeSpawner.apply_user_options static method.""" + """Tests for BeakerKubeSpawner.apply_user_options.""" def test_sets_beaker_context_from_user_options(self): """apply_user_options should set beaker_context from contextSlug.""" - spawner = MagicMock() - spawner.beaker_context = "" - spawner.context_config = {} + spawner = spawner_for_apply_user_options() user_options = {"contextSlug": "weather-context"} - BeakerKubeSpawner.apply_user_options(spawner, user_options) + BeakerKubeSpawner.apply_user_options(spawner, spawner, user_options) assert spawner.beaker_context == "weather-context" def test_sets_beaker_context_strips_prefix(self): """apply_user_options should strip prefix from contextSlug if it contains ':'.""" - spawner = MagicMock() - spawner.beaker_context = "" - spawner.context_config = {} + spawner = spawner_for_apply_user_options() user_options = {"contextSlug": "pkg:weather-context"} - BeakerKubeSpawner.apply_user_options(spawner, user_options) + BeakerKubeSpawner.apply_user_options(spawner, spawner, user_options) assert spawner.beaker_context == "weather-context" def test_sets_context_config_from_user_options(self): """apply_user_options should set context_config from contextOptions.""" - spawner = MagicMock() - spawner.beaker_context = "" - spawner.context_config = {} + spawner = spawner_for_apply_user_options() user_options = { "contextOptions": {"key1": "value1", "key2": "value2"} } - BeakerKubeSpawner.apply_user_options(spawner, user_options) + BeakerKubeSpawner.apply_user_options(spawner, spawner, user_options) assert spawner.context_config == {"key1": "value1", "key2": "value2"} def test_handles_missing_context_option(self): """apply_user_options should handle missing contextSlug gracefully.""" - spawner = MagicMock() - spawner.beaker_context = "original" - spawner.context_config = {} + spawner = spawner_for_apply_user_options(beaker_context="original") user_options = {} - BeakerKubeSpawner.apply_user_options(spawner, user_options) + BeakerKubeSpawner.apply_user_options(spawner, spawner, user_options) # beaker_context should not be modified assert spawner.beaker_context == "original" def test_handles_missing_context_options(self): """apply_user_options should handle missing contextOptions gracefully.""" - spawner = MagicMock() - spawner.beaker_context = "" - spawner.context_config = {"original": "config"} + spawner = spawner_for_apply_user_options( + context_config={"original": "config"} + ) user_options = {"contextSlug": "new-context"} - BeakerKubeSpawner.apply_user_options(spawner, user_options) + BeakerKubeSpawner.apply_user_options(spawner, spawner, user_options) # context_config should not be modified assert spawner.context_config == {"original": "config"} def test_handles_empty_user_options(self): """apply_user_options should handle empty options dict.""" - spawner = MagicMock() - spawner.beaker_context = "original" - spawner.context_config = {"original": "config"} + spawner = spawner_for_apply_user_options( + beaker_context="original", + context_config={"original": "config"}, + ) user_options = {} # Should not raise - BeakerKubeSpawner.apply_user_options(spawner, user_options) + BeakerKubeSpawner.apply_user_options(spawner, spawner, user_options) assert spawner.beaker_context == "original" assert spawner.context_config == {"original": "config"} @@ -236,14 +260,7 @@ def db(self): @pytest.fixture def spawner(self, db): - spawner = MagicMock() - spawner.db = db - spawner.debug = False - spawner.default_tag = "latest" - spawner.extra_container_config = {} - spawner.node_env = {} - spawner.node_policy_overrides = {} - return spawner + return spawner_for_apply_user_options(db=db) def test_passes_only_secrets_with_overrides(self, db, spawner): """Secrets at their defaults send nothing; the node already assumes defaults.""" @@ -259,7 +276,7 @@ def test_passes_only_secrets_with_overrides(self, db, spawner): ]) db.commit() - BeakerKubeSpawner.apply_user_options(spawner, {}) + BeakerKubeSpawner.apply_user_options(spawner, spawner, {}) # Both values reach the pod... assert spawner.node_env == {"PRIVATE_LLM_KEY": "a", "SHARED_API_KEY": "b"} @@ -290,7 +307,9 @@ def test_node_secret_policies_override_global(self, db, spawner): ]) db.commit() - BeakerKubeSpawner.apply_user_options(spawner, {"nodeSlug": "test-node"}) + BeakerKubeSpawner.apply_user_options( + spawner, spawner, {"nodeSlug": "test-node"} + ) assert spawner.node_env["SHARED_API_KEY"] == "node" assert spawner.node_policy_overrides == {"SHARED_API_KEY": {"ui_message_policy": "last4"}} @@ -312,7 +331,9 @@ def test_node_secret_without_policies_clears_inherited_ones(self, db, spawner): ]) db.commit() - BeakerKubeSpawner.apply_user_options(spawner, {"nodeSlug": "test-node"}) + BeakerKubeSpawner.apply_user_options( + spawner, spawner, {"nodeSlug": "test-node"} + ) assert spawner.node_env["SHARED_API_KEY"] == "node" assert spawner.node_policy_overrides == {} @@ -336,7 +357,9 @@ def test_context_disabled_secret_sends_no_policies(self, db, spawner): )) db.commit() - BeakerKubeSpawner.apply_user_options(spawner, {"contextSlug": "weather"}) + BeakerKubeSpawner.apply_user_options( + spawner, spawner, {"contextSlug": "weather"} + ) assert spawner.node_env == {} assert spawner.node_policy_overrides == {} diff --git a/tests/unit/services/test_aws_ecs_dashboard.py b/tests/unit/services/test_aws_ecs_dashboard.py new file mode 100644 index 0000000..8c6d8c6 --- /dev/null +++ b/tests/unit/services/test_aws_ecs_dashboard.py @@ -0,0 +1,175 @@ +"""Tests for the ECS dashboard service.""" + +from unittest.mock import MagicMock + +import boto3 + +from traitlets.config import LoggingConfigurable + +from beakerhub.runtimes.aws_ecs import AwsEcsRuntime +from beakerhub.services.dashboard.aws_ecs_dashboard import AwsEcsDashboardService + + +def make_service(): + ecs_client = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + logs_client = boto3.client( + "logs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + runtime = AwsEcsRuntime( + aws_region="us-east-1", + cluster_name="beakerhub", + log_group="/beakerhub/sessions", + log_stream_prefix="beakerhub", + ecs_client=ecs_client, + logs_client=logs_client, + ) + for method in ( + "describe_clusters", "list_services", "describe_services", "list_tasks", + "describe_tasks", "list_container_instances", "describe_container_instances", + ): + setattr(runtime.ecs_client, method, MagicMock()) + runtime.logs_client.get_log_events = MagicMock() + parent = LoggingConfigurable() + parent.runtime = runtime + return AwsEcsDashboardService(parent=parent), runtime + + +def test_dashboard_returns_provider_neutral_runtime_information(): + service, runtime = make_service() + runtime.ecs_client.describe_clusters.return_value = { + "clusters": [{"clusterArn": "arn:aws:ecs:region:account:cluster/beakerhub"}] + } + runtime.ecs_client.list_services.return_value = {"serviceArns": ["service-arn"]} + runtime.ecs_client.describe_services.return_value = { + "services": [ + { + "serviceName": "sessions", + "status": "ACTIVE", + "runningCount": 2, + "desiredCount": 2, + } + ] + } + runtime.ecs_client.list_tasks.return_value = {"taskArns": ["task-arn"]} + runtime.ecs_client.describe_tasks.return_value = { + "tasks": [ + { + "taskArn": "task-arn", + "group": "service:sessions", + "lastStatus": "RUNNING", + } + ] + } + runtime.ecs_client.list_container_instances.return_value = { + "containerInstanceArns": [] + } + + result = service.get_dashboard() + + assert result["available"] is True + assert result["runtime"] == {"provider": "Amazon ECS", "scope": "beakerhub"} + assert result["summary"][0]["label"] == "Services" + assert result["summary"][2]["detail"] == ( + "Compute is managed by AWS Fargate. This cluster has no registered " + "container instances." + ) + assert result["resources_empty_message"] == ( + "AWS Fargate manages the underlying compute resources. Per-instance " + "capacity is not available." + ) + assert result["workloads"][0]["name"] == "sessions" + assert result["workloads"][0]["kind"] == "Service" + assert result["workloads"][0]["status"] == "active" + assert result["workloads"][0]["detail"] == "2 running / 2 desired" + assert result["workloads"][0]["children"] == [ + { + "name": "task-arn", + "kind": "Task", + "status": "running", + "detail": "", + } + ] + + +def test_dashboard_reports_recent_stopped_tasks_that_failed(): + service, runtime = make_service() + runtime.ecs_client.describe_clusters.return_value = { + "clusters": [{"clusterArn": "arn:aws:ecs:region:account:cluster/beakerhub"}] + } + runtime.ecs_client.list_services.return_value = {"serviceArns": []} + runtime.ecs_client.list_tasks.side_effect = [ + {"taskArns": []}, + {"taskArns": ["failed-task"]}, + ] + runtime.ecs_client.describe_tasks.return_value = { + "tasks": [ + { + "taskArn": "failed-task", + "lastStatus": "STOPPED", + "stoppedReason": "Essential container exited", + "containers": [{"exitCode": 1}], + } + ] + } + runtime.ecs_client.list_container_instances.return_value = { + "containerInstanceArns": [] + } + + result = service.get_dashboard() + + assert result["summary"][1]["detail"] == "0 running, 0 pending, 1 failed" + assert result["summary"][1]["severity"] == "danger" + assert result["alerts"][0]["object"] == "failed-task" + + +def test_dashboard_does_not_report_user_stopped_tasks_as_failures(): + service, _ = make_service() + + assert service._failed_tasks( + [{"stopCode": "UserInitiated", "containers": [{"exitCode": 137}]}] + ) == [] + + +def test_dashboard_reports_task_start_failures_without_container_exit_codes(): + service, _ = make_service() + + assert service._failed_tasks([{"stopCode": "TaskFailedToStart", "containers": []}]) + + +def test_session_logs_uses_the_task_container_when_notebook_is_requested(): + service, runtime = make_service() + task_arn = "arn:aws:ecs:region:account:task/beakerhub/task-id" + runtime.ecs_client.describe_tasks.return_value = { + "tasks": [ + { + "taskArn": task_arn, + "containers": [{"name": "task-container"}], + } + ] + } + runtime.logs_client.get_log_events.side_effect = [ + {"events": [{"message": "second"}], "nextBackwardToken": "previous"}, + {"events": [{"message": "first"}], "nextBackwardToken": "previous"}, + ] + + result = service.get_session_logs(task_arn, "notebook", 5000) + + assert result["runtime_name"] == "task-id" + assert result["container"] == "task-container" + assert result["logs"] == "first\nsecond" + assert result["truncated"] is False + runtime.logs_client.get_log_events.assert_called_with( + logGroupName="/beakerhub/sessions", + logStreamName="beakerhub/task-container/task-id", + startFromHead=False, + limit=4999, + nextToken="previous", + ) diff --git a/tests/unit/services/test_aws_ecs_task_runner.py b/tests/unit/services/test_aws_ecs_task_runner.py new file mode 100644 index 0000000..8241b89 --- /dev/null +++ b/tests/unit/services/test_aws_ecs_task_runner.py @@ -0,0 +1,127 @@ +"""Tests for the ECS task-runner runtime adapter.""" + +from dataclasses import dataclass +from unittest.mock import Mock + +import boto3 + +from beakerhub.runtimes.aws_ecs import AwsEcsRuntime +from beakerhub.services.task.aws_ecs_task_runner import AwsEcsTaskRunnerService +from beakerhub.tasks.base import BaseImageTask + + +@dataclass(frozen=True) +class ExampleImageTask(BaseImageTask): + @property + def task_type(self) -> str: + return "example" + + +def runtime(**overrides) -> AwsEcsRuntime: + ecs = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + logs = boto3.client( + "logs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + values = { + "ecs_client": ecs, + "logs_client": logs, + "cluster_name": "runtime-cluster", + "subnets": ["subnet-runtime"], + "security_groups": ["sg-runtime"], + "log_group": "/beakerhub/runtime", + "execution_role_arn": "arn:aws:iam::123456789012:role/execution", + "task_role_arn": "arn:aws:iam::123456789012:role/task", + "launch_options": {"enableExecuteCommand": True}, + } + values.update(overrides) + return AwsEcsRuntime(**values) + + +def test_runner_defaults_to_the_shared_runtime_configuration(): + shared_runtime = runtime() + runner = AwsEcsTaskRunnerService(runtime=shared_runtime) + + assert runner.runtime is shared_runtime + assert runner.cluster_name == "runtime-cluster" + assert runner.subnets == ["subnet-runtime"] + assert runner.security_groups == ["sg-runtime"] + assert runner.log_group == "/beakerhub/runtime" + assert runner.execution_role_arn.endswith("role/execution") + assert runner.launch_options == {"enableExecuteCommand": True} + + +def test_submit_builds_a_runtime_definition_and_launches_the_process(): + shared_runtime = runtime() + shared_runtime.ecs_client.list_task_definitions = Mock( + return_value={"taskDefinitionArns": []} + ) + shared_runtime.ecs_client.register_task_definition = Mock( + return_value={"taskDefinition": {"taskDefinitionArn": "definition-arn"}} + ) + shared_runtime.ecs_client.run_task = Mock( + return_value={"tasks": [{"taskArn": "task-arn"}]} + ) + runner = AwsEcsTaskRunnerService( + runtime=shared_runtime, + security_groups=["sg-task"], + assign_public_ip=True, + ) + + process = runner.submit( + ExampleImageTask( + image="example:latest", + environment={"FROM_TASK": "value"}, + ) + ) + + assert process.runtime is shared_runtime + assert process.external_id == "task-arn" + definition_request = shared_runtime.ecs_client.register_task_definition.call_args.kwargs + assert definition_request["executionRoleArn"].endswith("role/execution") + assert definition_request["taskRoleArn"].endswith("role/task") + task_request = shared_runtime.ecs_client.run_task.call_args.kwargs + assert task_request["cluster"] == "runtime-cluster" + assert task_request["networkConfiguration"] == { + "awsvpcConfiguration": { + "subnets": ["subnet-runtime"], + "securityGroups": ["sg-task"], + "assignPublicIp": "ENABLED", + } + } + assert task_request["overrides"]["containerOverrides"] == [ + { + "name": "task-container", + "environment": [{"name": "FROM_TASK", "value": "value"}], + } + ] + assert {tag["key"]: tag["value"] for tag in task_request["tags"]} == { + "beakerhub/task-type": "example", + "beakerhub-process": "true", + "beakerhub-process-type": "task", + } + + +def test_runner_overrides_are_used_when_reconstructing_a_process(): + shared_runtime = runtime() + runner = AwsEcsTaskRunnerService( + runtime=shared_runtime, + cluster_name="task-cluster", + log_group="/beakerhub/tasks", + task_overrides={"cpu": "512"}, + ) + + process = runner.get_process("task-arn") + + assert process.runtime is shared_runtime + assert process.external_id == "task-arn" + assert process.definition.cluster_name == "task-cluster" + assert process.definition.log_group == "/beakerhub/tasks" + assert process.definition.task_overrides == {"cpu": "512"} diff --git a/tests/unit/services/test_dashboard_handlers.py b/tests/unit/services/test_dashboard_handlers.py new file mode 100644 index 0000000..75f78f2 --- /dev/null +++ b/tests/unit/services/test_dashboard_handlers.py @@ -0,0 +1,36 @@ +"""Tests for dashboard handler data enrichment.""" + +from unittest.mock import MagicMock + +from beakerhub.services.dashboard.handlers import AdminDashboardClusterHandler + + +def test_dashboard_task_rows_include_the_owned_session(): + spawner = MagicMock() + spawner.state = { + "task_arn": "arn:aws:ecs:us-east-1:123456789:task/cluster/task-id" + } + spawner.user.name = "ada" + spawner.name = "analysis" + + handler = MagicMock() + handler.db.query.return_value.join.return_value.filter.return_value.all.return_value = [ + spawner + ] + dashboard = { + "workloads": [ + { + "name": "notebook-service", + "kind": "Service", + "children": [ + {"name": "task-id", "kind": "Task"}, + ], + } + ] + } + + AdminDashboardClusterHandler._add_task_sessions(handler, dashboard) + + assert dashboard["workloads"][0]["children"][0]["sessions"] == [ + {"user": "ada", "name": "analysis"} + ] diff --git a/tests/unit/services/test_kubernetes_task_runner.py b/tests/unit/services/test_kubernetes_task_runner.py new file mode 100644 index 0000000..54aa537 --- /dev/null +++ b/tests/unit/services/test_kubernetes_task_runner.py @@ -0,0 +1,62 @@ +"""Tests for the Kubernetes task-runner runtime adapter.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +from beakerhub.runtimes.kubernetes import KubernetesRuntime +from beakerhub.services.task.kubernetes_task_runner import KubernetesTaskRunnerService +from beakerhub.tasks.image_import.task import ImageImportTask + + +def runtime(**overrides): + values = { + "namespace": "runtime-namespace", + "node_selector": {"pool": "runtime"}, + "tolerations": [{"key": "runtime", "operator": "Exists"}], + } + values.update(overrides) + instance = KubernetesRuntime(**values) + instance.batch_api = Mock() + instance.core_api = Mock() + return instance + + +def image_import_task(): + return ImageImportTask.from_node_image( + SimpleNamespace(default_img_string="registry.example/node:latest") + ) + + +def test_submit_uses_the_shared_runtime_defaults(): + shared_runtime = runtime() + runner = KubernetesTaskRunnerService(runtime=shared_runtime) + + process = runner.submit(image_import_task()) + + request = shared_runtime.batch_api.create_namespaced_job.call_args.kwargs + pod_spec = request["body"].spec.template.spec + assert process.runtime is shared_runtime + assert request["namespace"] == "runtime-namespace" + assert pod_spec.node_selector == {"pool": "runtime"} + assert pod_spec.tolerations[0].key == "runtime" + assert request["body"].metadata.labels["beakerhub/task-type"] == "context_import" + + +def test_runner_configuration_overrides_runtime_defaults_and_reconstructs_processes(): + shared_runtime = runtime() + runner = KubernetesTaskRunnerService( + runtime=shared_runtime, + namespace="task-namespace", + node_selector={"pool": "tasks"}, + tolerations=[{"key": "tasks", "operator": "Exists"}], + ) + + runner.submit(image_import_task()) + submitted = shared_runtime.batch_api.create_namespaced_job.call_args.kwargs["body"] + restored = runner.get_process("beaker-task-context-import-123") + + assert submitted.metadata.namespace == "task-namespace" + assert submitted.spec.template.spec.node_selector == {"pool": "tasks"} + assert submitted.spec.template.spec.tolerations[0].key == "tasks" + assert restored.runtime is shared_runtime + assert restored.definition.namespace == "task-namespace" diff --git a/tests/unit/services/test_service_wiring.py b/tests/unit/services/test_service_wiring.py new file mode 100644 index 0000000..e978b94 --- /dev/null +++ b/tests/unit/services/test_service_wiring.py @@ -0,0 +1,147 @@ +"""Tests for BeakerHub service selection and task delegation.""" + +from unittest.mock import MagicMock + +from traitlets.config import Config + +from beakerhub.app import BeakerHub +from beakerhub.runtimes.aws_ecs import AwsEcsRuntime, AwsEcsRuntimeBundle +from beakerhub.runtimes.kubernetes import KubernetesRuntime, KubernetesRuntimeBundle +from beakerhub.services.dashboard.aws_ecs_dashboard import AwsEcsDashboardService +from beakerhub.services.dashboard.base import BaseDashboardService +from beakerhub.services.dashboard.kubernetes_dashboard import KubernetesDashboardService +from beakerhub.services.task.aws_ecs_task_runner import AwsEcsTaskRunnerService +from beakerhub.services.task.kubernetes_task_runner import KubernetesTaskRunnerService +from beakerhub.services.task.base import BaseTaskRunnerService +from beakerhub.tasks.image_import.task import ImageImportTask, launch_import_task + + +class DummyTaskRunnerService(BaseTaskRunnerService): + """Task runner used to verify application service selection.""" + + +class DummyDashboardService(BaseDashboardService): + """Dashboard service used to verify application service selection.""" + + +def test_application_instantiates_configured_services(): + config = Config() + config.BeakerHub.task_runner_class = DummyTaskRunnerService + config.BeakerHub.dashboard_service_class = DummyDashboardService + + app = BeakerHub(config=config) + + assert isinstance(app.task_runner, DummyTaskRunnerService) + assert app.task_runner.parent is app + assert isinstance(app.dashboard_service, DummyDashboardService) + assert app.dashboard_service.parent is app + + +def test_application_registers_service_classes(): + app = BeakerHub() + + assert BaseTaskRunnerService in app.classes + assert BaseDashboardService in app.classes + + +def test_runtime_bundle_selects_services_with_application_owned_runtime(): + config = Config() + config.BeakerHub.runtime_bundle_class = KubernetesRuntimeBundle + config.KubernetesRuntime.namespace = "configured-runtime" + config.KubernetesRuntime.node_selector = {"pool": "compute"} + config.KubernetesRuntime.tolerations = [{"key": "dedicated", "operator": "Exists"}] + config.KubernetesRuntime.service_account = "runtime-workload" + config.KubernetesRuntime.base_labels = {"environment": "test"} + + app = BeakerHub(config=config) + + assert isinstance(app.runtime_bundle, KubernetesRuntimeBundle) + assert app.runtime.parent is app + assert app.runtime.namespace == "configured-runtime" + assert app.runtime.node_selector == {"pool": "compute"} + assert app.runtime.tolerations == [{"key": "dedicated", "operator": "Exists"}] + assert app.runtime.service_account == "runtime-workload" + assert app.runtime.base_labels == {"environment": "test"} + assert isinstance(app.task_runner, KubernetesTaskRunnerService) + assert app.task_runner.runtime is app.runtime + assert isinstance(app.dashboard_service, KubernetesDashboardService) + assert app.dashboard_service._runtime() is app.runtime + assert app.spawner_class.__name__ == "BeakerKubeSpawner" + + +def test_ecs_runtime_bundle_selects_ecs_services_and_runtime_values(): + config = Config() + config.BeakerHub.runtime_bundle_class = AwsEcsRuntimeBundle + config.AwsEcsRuntime.aws_region = "us-west-2" + config.AwsEcsRuntime.cluster_name = "configured-cluster" + config.AwsEcsRuntime.subnets = ["subnet-123"] + config.AwsEcsRuntime.log_group = "/beakerhub/ecs" + + app = BeakerHub(config=config) + + assert isinstance(app.runtime, AwsEcsRuntime) + assert app.runtime.parent is app + assert app.runtime.aws_region == "us-west-2" + assert app.runtime.cluster_name == "configured-cluster" + assert app.runtime.subnets == ["subnet-123"] + assert app.runtime.log_group == "/beakerhub/ecs" + assert isinstance(app.task_runner, AwsEcsTaskRunnerService) + assert app.task_runner.runtime is app.runtime + assert isinstance(app.dashboard_service, AwsEcsDashboardService) + assert app.spawner_class.__name__ == "BeakerAwsECSSpawner" + + +def test_runtime_bundle_runtime_receives_late_configuration_updates(): + config = Config() + config.BeakerHub.runtime_bundle_class = KubernetesRuntimeBundle + app = BeakerHub(config=config) + runtime = app.runtime + + update = Config() + update.KubernetesRuntime.namespace = "updated-runtime" + update.KubernetesRuntime.node_selector = {"pool": "updated"} + app.update_config(update) + + assert runtime.namespace == "updated-runtime" + assert runtime.node_selector == {"pool": "updated"} + + +def test_default_runtime_is_kubernetes_without_a_bundle(): + app = BeakerHub() + + assert isinstance(app.runtime, KubernetesRuntime) + assert app.runtime_bundle is None + + +def test_task_runner_receives_config_loaded_after_its_creation(): + app = BeakerHub() + runner = app.task_runner + config = Config() + config.KubernetesTaskRunnerService.node_image_resources = { + "limits": {"cpu": "1"} + } + + app.update_config(config) + + assert runner.node_image_resources == {"limits": {"cpu": "1"}} + + +def test_image_import_delegates_submission_to_task_runner(): + db = MagicMock() + db.query.return_value.filter.return_value.first.return_value = None + node_image = MagicMock(id=12, slug="example") + app = MagicMock(hub_connect_url="http://hub.internal") + app.task_runner.submit.return_value = MagicMock(external_id="external-task-id") + + task = launch_import_task(db, app, node_image) + + app.task_runner.submit.assert_called_once() + submitted_task = app.task_runner.submit.call_args.args[0] + assert isinstance(submitted_task, ImageImportTask) + assert submitted_task.node_image is node_image + assert submitted_task.image == node_image.default_img_string + assert submitted_task.entrypoint == ("sh", "-c") + assert submitted_task.command == ("beaker context dump",) + assert task.job_name == "external-task-id" + assert task.status == "running" + assert db.commit.call_count == 2 diff --git a/tests/unit/services/test_task_base.py b/tests/unit/services/test_task_base.py new file mode 100644 index 0000000..1b1d357 --- /dev/null +++ b/tests/unit/services/test_task_base.py @@ -0,0 +1,78 @@ +"""Tests for task-runner delegation to runtime processes.""" + +import pytest + +from beakerhub.runtimes.base import ( + BaseDefinition, + BaseProcess, + BaseRuntime, + ProcessOutput, + ProcessStatus, +) +from beakerhub.services.task.base import BaseTaskRunnerService, TaskOutput, TaskStatus + + +class CompleteProcess(BaseProcess): + def describe(self) -> ProcessStatus: + return ProcessStatus("completed") + + def collect_output(self) -> ProcessOutput: + return ProcessOutput("output", "") + + def stop(self) -> None: + return None + + +class PendingProcess(CompleteProcess): + def describe(self) -> ProcessStatus: + return ProcessStatus("running") + + +class ProcessTaskRunner(BaseTaskRunnerService): + def __init__(self, process: BaseProcess, **kwargs): + self.process = process + super().__init__(**kwargs) + + def get_process(self, external_id: str) -> BaseProcess: + assert external_id == self.process.external_id + return self.process + + +def complete_process() -> CompleteProcess: + return CompleteProcess( + BaseDefinition(), + runtime=BaseRuntime(), + external_id="runtime-123", + ) + + +def test_task_status_and_output_are_runtime_contract_aliases(): + assert TaskStatus is ProcessStatus + assert TaskOutput is ProcessOutput + + +def test_task_runner_delegates_persisted_identifier_to_process(): + process = complete_process() + runner = ProcessTaskRunner(process) + + assert runner.get_status("runtime-123") == ProcessStatus("completed") + assert runner.get_output("runtime-123") == ProcessOutput("output", "") + + +@pytest.mark.asyncio +async def test_process_returns_terminal_status(): + process = complete_process() + + assert await process.await_completion() == ProcessStatus("completed") + + +@pytest.mark.asyncio +async def test_process_raises_on_timeout(): + process = PendingProcess( + BaseDefinition(), + runtime=BaseRuntime(), + external_id="runtime-123", + ) + + with pytest.raises(TimeoutError, match="runtime-123"): + await process.await_completion(timeout=0) diff --git a/tests/unit/spawner/__init__.py b/tests/unit/spawner/__init__.py deleted file mode 100644 index 805c7dc..0000000 --- a/tests/unit/spawner/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: 2024-present Jataware Corp -# -# SPDX-License-Identifier: MIT diff --git a/tests/unit/tasks/test_image_import.py b/tests/unit/tasks/test_image_import.py new file mode 100644 index 0000000..af54645 --- /dev/null +++ b/tests/unit/tasks/test_image_import.py @@ -0,0 +1,55 @@ +"""Tests for node-image import task definition behavior.""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from beakerhub.runtimes.base import ProcessOutput, ProcessStatus +from beakerhub.tasks.image_import.task import ImageImportTask + + +def node_image(): + return SimpleNamespace(default_img_string="registry.example/node:latest") + + +def test_from_node_image_builds_the_standard_context_dump_workload(): + task = ImageImportTask.from_node_image(node_image()) + + assert task.image == "registry.example/node:latest" + assert task.entrypoint == ("sh", "-c") + assert task.command == ("beaker context dump",) + assert task.task_type == "context_import" + + +def test_on_success_ingests_output_into_the_persisted_task(): + node = node_image() + definition = ImageImportTask.from_node_image(node) + persisted_task = SimpleNamespace(result=None) + db = Mock() + output = ProcessOutput('{"contexts": []}', "") + + with ( + patch( + "beakerhub.tasks.image_import.task.object_session", return_value=db + ), + patch( + "beakerhub.tasks.image_import.task.ingest_import_output", + return_value={"contexts": 1}, + ) as ingest, + ): + definition.on_success(persisted_task, ProcessStatus("completed"), output) + + ingest.assert_called_once_with(db, node, output) + assert persisted_task.result == {"contexts": 1} + + +def test_on_failure_stores_runtime_diagnostics(): + definition = ImageImportTask.from_node_image(node_image()) + persisted_task = SimpleNamespace(error=None) + + definition.on_failure( + persisted_task, + ProcessStatus("failed", "Container exited unsuccessfully"), + ProcessOutput("", "Import command failed"), + ) + + assert persisted_task.error == "Container exited unsuccessfully. Import command failed" diff --git a/ui/src/components/admin/PodLogViewer.vue b/ui/src/components/admin/SessionLogViewer.vue similarity index 92% rename from ui/src/components/admin/PodLogViewer.vue rename to ui/src/components/admin/SessionLogViewer.vue index 9c18bcc..373778b 100644 --- a/ui/src/components/admin/PodLogViewer.vue +++ b/ui/src/components/admin/SessionLogViewer.vue @@ -33,8 +33,8 @@
- - {{ podName }} + + {{ runtimeName }} Fetched {{ formatTime(timestamp) }} @@ -78,7 +78,7 @@ const loading = ref(false); const loadingMore = ref(false); const logs = ref(null); const error = ref(null); -const podName = ref(null); +const runtimeName = ref(null); const timestamp = ref(null); const truncated = ref(false); const currentTailLines = ref(5000); @@ -88,7 +88,7 @@ function onShow() { currentTailLines.value = 5000; logs.value = null; error.value = null; - podName.value = null; + runtimeName.value = null; timestamp.value = null; truncated.value = false; fetchLogs(); @@ -103,16 +103,15 @@ async function fetchLogs() { loading.value = true; error.value = null; try { - const result = await adminStore.fetchPodLogs( + const result = await adminStore.fetchSessionLogs( props.username, props.serverName, currentTailLines.value ); logs.value = result.logs; - podName.value = result.pod_name; + runtimeName.value = result.runtime_name; timestamp.value = result.timestamp; truncated.value = result.truncated; - // Scroll to bottom after render requestAnimationFrame(() => { scrollToBottom(); }); @@ -134,16 +133,15 @@ async function loadMore() { currentTailLines.value *= 2; try { - const result = await adminStore.fetchPodLogs( + const result = await adminStore.fetchSessionLogs( props.username, props.serverName, currentTailLines.value ); logs.value = result.logs; - podName.value = result.pod_name; + runtimeName.value = result.runtime_name; timestamp.value = result.timestamp; truncated.value = result.truncated; - // Preserve scroll position relative to existing content requestAnimationFrame(() => { const newScrollHeight = container.scrollHeight; container.scrollTop = newScrollHeight - previousScrollHeight; @@ -240,7 +238,7 @@ function formatTime(isoString: string): string { flex-shrink: 0; } -.pod-name { +.runtime-name { font-family: monospace; font-size: 0.8rem; color: var(--p-text-secondary-color); diff --git a/ui/src/pages/Dashboard.vue b/ui/src/pages/Dashboard.vue index 5e3fcf6..e13a6a7 100644 --- a/ui/src/pages/Dashboard.vue +++ b/ui/src/pages/Dashboard.vue @@ -487,7 +487,7 @@ function getStatusSeverity(session: SessionListItem): "success" | "warn" | "info function getStatusLabel(session: SessionListItem): string { if (session.stopped) return 'Stopped'; if (session.pending === 'spawn') return 'Starting...'; - if (session.pending === 'stop') return 'Stopping...'; + if (session.pending === 'stop') return 'Shutting Down...'; if (session.ready) return 'Running'; return 'Unknown'; } diff --git a/ui/src/pages/admin/AdminDashboard.vue b/ui/src/pages/admin/AdminDashboard.vue index 0113a32..01d6846 100644 --- a/ui/src/pages/admin/AdminDashboard.vue +++ b/ui/src/pages/admin/AdminDashboard.vue @@ -117,7 +117,7 @@ @@ -194,188 +194,106 @@
- +
-

Kubernetes Cluster

- {{ clusterInfo.namespace }} +
+

Cluster Information

+ + {{ clusterInfo.runtime.provider }} · {{ clusterInfo.runtime.scope }} + +
- Cluster information is unavailable. {{ clusterInfo.error || 'Insufficient permissions or not running in-cluster.' }} + Runtime information is unavailable. {{ clusterInfo.error || 'The service did not return status information.' }}
- ({ user: '', name: '' }); const counts = computed(() => adminStore.dashboardSummary?.counts ?? null); const recentImports = computed(() => adminStore.dashboardSummary?.recent_imports ?? []); const clusterInfo = computed(() => adminStore.clusterInfo); -const clusterNodes = computed(() => adminStore.clusterInfo?.nodes ?? []); -const helmReleases = computed(() => adminStore.clusterInfo?.helm_releases ?? []); -const appVersion = computed(() => adminStore.dashboardSummary?.app_version ?? '...'); const serverList = computed(() => { const rows: ServerRow[] = []; @@ -474,6 +389,10 @@ async function refreshAll(silent = false) { } } +function pendingStatus(status: string): string { + return status === 'stop' ? 'Shutting Down' : status; +} + function openLogs(server: ServerRow) { logsTarget.value = { user: server.user, name: server.name }; logsDialogVisible.value = true; @@ -537,78 +456,25 @@ function formatImportResult(result: Record): string { return parts.join(', '); } -function parseCpuToMillicores(value: string | null | undefined): number { - if (!value || value === '?') return 0; - if (value.endsWith('m')) return parseInt(value.slice(0, -1), 10) || 0; - return (parseFloat(value) || 0) * 1000; -} - -function parseMemoryToKi(value: string | null | undefined): number { - if (!value || value === '?') return 0; - const match = value.match(/^([\d.]+)(Ki|Mi|Gi|Ti)?$/); - if (!match) return 0; - const num = parseFloat(match[1]); - switch (match[2]) { - case 'Ti': return num * 1024 * 1024 * 1024; - case 'Gi': return num * 1024 * 1024; - case 'Mi': return num * 1024; - case 'Ki': return num; - default: return num / 1024; - } -} - -function utilizationPct(allocated: string | null | undefined, allocatable: string | null | undefined): number { - if (!allocated || !allocatable || allocatable === '?' || allocated === '?') return 0; - let used: number; - let total: number; - // Memory values have Ki/Mi/Gi/Ti suffixes - if (allocated.match(/[KMGTi]i?$/) || allocatable.match(/[KMGTi]i?$/)) { - used = parseMemoryToKi(allocated); - total = parseMemoryToKi(allocatable); - } else { - // CPU (e.g. '500m', '2') or plain numbers (pods) - used = parseCpuToMillicores(allocated); - total = parseCpuToMillicores(allocatable); - } - if (total === 0) return 0; - return Math.min(Math.round((used / total) * 100), 100); -} - -function utilizationClass(allocated: string | null | undefined, allocatable: string | null | undefined): string { - const pct = utilizationPct(allocated, allocatable); - if (pct >= 90) return 'bar-danger'; - if (pct >= 70) return 'bar-warn'; - return 'bar-ok'; -} - -function formatMemory(value: string | null | undefined): string { - if (!value) return '?'; - // K8s memory is typically in Ki (kibibytes) - const match = value.match(/^(\d+)(Ki|Mi|Gi|Ti)?$/); - if (!match) return value; - const num = parseInt(match[1], 10); - const unit = match[2] || ''; - if (unit === 'Ki') { - if (num >= 1048576) return `${(num / 1048576).toFixed(1)}Ti`; - if (num >= 1024) return `${(num / 1024).toFixed(1)}Gi`; - return `${num}Ki`; - } - if (unit === 'Mi') { - if (num >= 1024) return `${(num / 1024).toFixed(1)}Gi`; - return `${num}Mi`; +function statusSeverity(status: string): "success" | "info" | "warn" | "danger" | "secondary" { + switch (status.toLowerCase()) { + case 'running': + case 'ready': + case 'active': + case 'bound': + return 'success'; + case 'pending': + case 'provisioning': + return 'warn'; + case 'failed': + case 'unavailable': + case 'stopped': + return 'danger'; + default: + return 'secondary'; } - return value; } -function phaseSeverity(phase: string): "success" | "info" | "warn" | "danger" | "secondary" { - switch (phase) { - case 'Running': return 'success'; - case 'Succeeded': return 'info'; - case 'Pending': return 'warn'; - case 'Failed': return 'danger'; - default: return 'secondary'; - } -}