From 48cc941ffb45395965a471e2f0a26bbe53723d88 Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Fri, 21 Aug 2026 10:02:47 -0600 Subject: [PATCH 1/8] Initial setup of services based infrastructure --- .../beakerhub/templates/beakerhub-config.yaml | 21 +- helm/beakerhub/templates/configmap.yaml | 2 +- implicit-kubernetes-backend-assumptions.md | 644 ++++++++++++++++++ pyproject.toml | 2 +- src/beakerhub/admin_handlers.py | 2 +- src/beakerhub/app.py | 120 ++-- src/beakerhub/cli/main.py | 2 +- src/beakerhub/nodes/__init__.py | 0 src/beakerhub/nodes/tasks.py | 250 ------- src/beakerhub/services/dashboard/__init__.py | 5 + .../services/dashboard/aws_ecs_dashboard.py | 20 + src/beakerhub/services/dashboard/base.py | 31 + src/beakerhub/services/dashboard/handlers.py | 202 ++++++ .../dashboard/kubernetes_dashboard.py} | 211 +----- .../services/periodic_tasks/__init__.py | 1 + .../idle_culler/__init__.py | 2 +- .../idle_culler/service.py} | 0 src/beakerhub/services/spawner/__init__.py | 5 + .../services/spawner/aws_ecs_spawner.py | 107 +++ src/beakerhub/services/spawner/base.py | 148 ++++ .../spawner/kubernetes_spawner.py} | 89 +-- .../{ => services}/spawner/provisioner.py | 4 +- src/beakerhub/services/task/__init__.py | 5 + .../services/task/aws_ecs_task_runner.py | 26 + src/beakerhub/services/task/base.py | 29 + .../task/handlers.py} | 116 +--- .../services/task/kubernetes_task_runner.py | 264 +++++++ src/beakerhub/spawner/__init__.py | 0 src/beakerhub/spawner/base.py | 39 -- src/beakerhub/tasks/__init__.py | 1 + src/beakerhub/tasks/image_import/__init__.py | 5 + src/beakerhub/tasks/image_import/task.py | 78 +++ tests/unit/services/spawner/__init__.py | 1 + .../unit/{ => services}/spawner/test_base.py | 4 +- .../{ => services}/spawner/test_kubernetes.py | 4 +- tests/unit/services/test_service_wiring.py | 64 ++ tests/unit/spawner/__init__.py | 3 - 37 files changed, 1772 insertions(+), 735 deletions(-) create mode 100644 implicit-kubernetes-backend-assumptions.md delete mode 100644 src/beakerhub/nodes/__init__.py delete mode 100644 src/beakerhub/nodes/tasks.py create mode 100644 src/beakerhub/services/dashboard/__init__.py create mode 100644 src/beakerhub/services/dashboard/aws_ecs_dashboard.py create mode 100644 src/beakerhub/services/dashboard/base.py create mode 100644 src/beakerhub/services/dashboard/handlers.py rename src/beakerhub/{dashboard_handlers.py => services/dashboard/kubernetes_dashboard.py} (68%) create mode 100644 src/beakerhub/services/periodic_tasks/__init__.py rename src/beakerhub/services/{ => periodic_tasks}/idle_culler/__init__.py (89%) rename src/beakerhub/services/{idle_culler/idle_culler_service.py => periodic_tasks/idle_culler/service.py} (100%) create mode 100644 src/beakerhub/services/spawner/__init__.py create mode 100644 src/beakerhub/services/spawner/aws_ecs_spawner.py create mode 100644 src/beakerhub/services/spawner/base.py rename src/beakerhub/{spawner/kubernetes.py => services/spawner/kubernetes_spawner.py} (57%) rename src/beakerhub/{ => services}/spawner/provisioner.py (61%) create mode 100644 src/beakerhub/services/task/__init__.py create mode 100644 src/beakerhub/services/task/aws_ecs_task_runner.py create mode 100644 src/beakerhub/services/task/base.py rename src/beakerhub/{nodes/import_handlers.py => services/task/handlers.py} (63%) create mode 100644 src/beakerhub/services/task/kubernetes_task_runner.py delete mode 100644 src/beakerhub/spawner/__init__.py delete mode 100644 src/beakerhub/spawner/base.py create mode 100644 src/beakerhub/tasks/__init__.py create mode 100644 src/beakerhub/tasks/image_import/__init__.py create mode 100644 src/beakerhub/tasks/image_import/task.py create mode 100644 tests/unit/services/spawner/__init__.py rename tests/unit/{ => services}/spawner/test_base.py (97%) rename tests/unit/{ => services}/spawner/test_kubernetes.py (99%) create mode 100644 tests/unit/services/test_service_wiring.py delete mode 100644 tests/unit/spawner/__init__.py diff --git a/helm/beakerhub/templates/beakerhub-config.yaml b/helm/beakerhub/templates/beakerhub-config.yaml index 79d5815..4f73d9b 100644 --- a/helm/beakerhub/templates/beakerhub-config.yaml +++ b/helm/beakerhub/templates/beakerhub-config.yaml @@ -191,19 +191,20 @@ 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.reporter_image = {{ printf "%s%s:%s" (default "" .Values.defaultRegistry) .Values.tasks.reporter.image.repository .Values.tasks.reporter.image.tag | quote }} +c.KubernetesTaskRunnerService.reporter_pull_policy = {{ .Values.tasks.reporter.image.pullPolicy | quote }} +c.KubernetesTaskRunnerService.reporter_resources = {{ .Values.tasks.reporter.resources | toJson }} +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..5c8f3ad 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,18 @@ 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.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 +37,25 @@ 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", + task_runner_class = Type( + KubernetesTaskRunnerService, + klass=BaseTaskRunnerService, config=True, - help="Full image reference for the task reporter container.", + help="Task-runner service used for background workloads.", ) - 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.", - ) - task_node_image_resources = Dict( - config=True, - help="Resource requests/limits for the node image init container in tasks.", + task_runner = Instance( + BaseTaskRunnerService, + allow_none=False, ) - task_backoff_limit = Integer( - 0, + dashboard_service_class = Type( + KubernetesDashboardService, + klass=BaseDashboardService, config=True, - help="Number of retries before marking a task Job as failed.", + help="Service used to collect runtime dashboard data and session logs.", ) - task_active_deadline_seconds = Integer( - 300, - config=True, - help="Maximum time in seconds a task Job can run before being terminated.", - ) - task_ttl_seconds_after_finished = Integer( - 600, - config=True, - help="Time in seconds to keep completed/failed Jobs before cleanup.", - ) - task_node_selector = Dict( - config=True, - help="Node selector for task pods.", - ) - task_tolerations = List( - config=True, - help="Tolerations for task pods.", - ) - 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( @@ -129,9 +112,31 @@ def _default_authenticator_class(self): @default("spawner_class") def _default_spawner_class(self): - from beakerhub.spawner.kubernetes import BeakerKubeSpawner + from beakerhub.services.spawner.kubernetes_spawner import BeakerKubeSpawner return BeakerKubeSpawner + @default("task_runner") + def _default_task_runner(self): + return self.task_runner_class(parent=self) + + @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 "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): return "beakerhub_config.py" @@ -154,8 +159,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 +185,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 +212,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/__init__.py b/src/beakerhub/nodes/__init__.py deleted file mode 100644 index e69de29..0000000 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/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..2ad72b4 --- /dev/null +++ b/src/beakerhub/services/dashboard/aws_ecs_dashboard.py @@ -0,0 +1,20 @@ +"""AWS ECS implementation skeleton for the dashboard service.""" + +from typing import Any + +from beakerhub.services.dashboard.base import BaseDashboardService + + +class AwsEcsDashboardService(BaseDashboardService): + """Provide AWS ECS runtime data to the BeakerHub dashboard.""" + + def get_dashboard(self) -> dict[str, Any]: + raise NotImplementedError("AWS ECS dashboard data is not implemented") + + def get_session_logs( + self, + session_id: str, + container: str, + tail_lines: int, + ) -> dict[str, Any]: + raise NotImplementedError("AWS ECS session logs are not implemented") diff --git a/src/beakerhub/services/dashboard/base.py b/src/beakerhub/services/dashboard/base.py new file mode 100644 index 0000000..e7d695f --- /dev/null +++ b/src/beakerhub/services/dashboard/base.py @@ -0,0 +1,31 @@ +"""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 provider-specific dashboard data.""" + 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..df1a38d --- /dev/null +++ b/src/beakerhub/services/dashboard/handlers.py @@ -0,0 +1,202 @@ +"""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() + 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)) + + +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)) + + try: + result = self.settings["app"].dashboard_service.get_session_logs( + session_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/pod-logs/([^/]+)/([^/]+)", + AdminSessionLogsHandler, + ), +] + +api_handlers = handlers diff --git a/src/beakerhub/dashboard_handlers.py b/src/beakerhub/services/dashboard/kubernetes_dashboard.py similarity index 68% rename from src/beakerhub/dashboard_handlers.py rename to src/beakerhub/services/dashboard/kubernetes_dashboard.py index f2e280a..eb121d8 100644 --- a/src/beakerhub/dashboard_handlers.py +++ b/src/beakerhub/services/dashboard/kubernetes_dashboard.py @@ -1,154 +1,28 @@ -"""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.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. - - Queries the K8s API for pod, PVC, event, and job information. - Returns available: false if K8s access is not available. - """ - - def compute_etag(self) -> None: - return None - - @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), - } +class KubernetesDashboardService(BaseDashboardService): + """Provide Kubernetes cluster data and Pod logs.""" - self.set_header("Content-Type", "application/json") - self.write(json.dumps(result)) + namespace = Unicode( + "beakerhub", + config=True, + help="Kubernetes namespace inspected by the dashboard service.", + ) - def _get_cluster_info(self) -> dict[str, Any]: + def get_dashboard(self) -> dict[str, Any]: from kubernetes import client as k8s_client from kubernetes import config as k8s_config @@ -163,8 +37,7 @@ def _get_cluster_info(self) -> dict[str, Any]: core_api = k8s_client.CoreV1Api() batch_api = k8s_client.BatchV1Api() - app = self.settings.get("app") - namespace = getattr(app, "task_namespace", "beakerhub") + namespace = self.namespace # Pods pods_result = self._get_pod_info(core_api, namespace) @@ -509,35 +382,7 @@ 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 @@ -549,12 +394,14 @@ def _get_pod_logs( try: k8s_config.load_kube_config() except k8s_config.ConfigException: - raise web.HTTPError(503, reason="No Kubernetes configuration found") + raise DashboardServiceError( + 503, + "No Kubernetes configuration found", + ) core_api = k8s_client.CoreV1Api() - 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 +409,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 +426,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 @@ -595,10 +443,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..1146626 --- /dev/null +++ b/src/beakerhub/services/spawner/__init__.py @@ -0,0 +1,5 @@ +"""Spawner services for BeakerHub session runtimes.""" + +from .base import BeakerSpawner, BeakerhubImageSpawner + +__all__ = ["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..7cab0ae --- /dev/null +++ b/src/beakerhub/services/spawner/aws_ecs_spawner.py @@ -0,0 +1,107 @@ +import json +from typing import cast, Optional + +import boto3 +from jupyterhub.spawner import Spawner +from kubespawner.spawner import KubeSpawner +from traitlets import default, validate, Unicode, Dict, List, Bool +from traitlets.config import Application + +from beakerhub.auth.user import BeakerhubUser +from beakerhub.services.spawner.base import BeakerSpawner, BeakerhubImageSpawner +from beakerhub.services.secrets import VAULT_ENV_VAR_LIST_KEY +from beakerhub import orm + + +class BeakerAwsECSSpawner(BeakerhubImageSpawner): + cluster_name: str = Unicode( + config=True, + ) + container_type: str = Unicode( + config=True, + ) + container_tags: dict[str, str] = Dict( + Unicode, + config=True, + ) + fargate: bool = Bool( + False, + config=True, + ) + + boto: boto3.client + # session_key: str + task_arn = Optional[str] + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.boto = boto3.client("ecs") + self.task_arn = None + + @property + def volume_configs(self) -> list[dict]: + return [] + + + def start(self): + overrides = {} + tasks = self.boto.run_task( + cluster=self.cluster_name, + count=1, + networkConfiguration={ + "assignPublicIp": "ENABLED", + }, + overrides=overrides, + tags=(self.container_tags + [{"key": "beaker-session", "value": self.session_id}]), + volumeConfigurations=self.volume_configs + ) + self.tasks = tasks + + def stop(self, now=False): + self.boto.stop_task( + cluster=self.cluster_name, + task=self.task_arn, + reason="Stopped by BeakerHub server" + ) + + + def poll(self): + """ + Check if the pod is still running. + + Uses the same interface as subprocess.Popen.poll(): if the pod is + still running, returns None. If the pod has exited, return the + exit code if we can determine it, or 1 if it has exited but we + don't know how. These are the return values JupyterHub expects. + + Note that a clean exit will have an exit code of zero, so it is + necessary to check that the returned value is None, rather than + just Falsy, to determine that the pod is still running. + """ + task_response = self.boto.describe_tasks( + cluster=self.cluster_name, + tasks=[self.task_arn], + ) + task = next((task for task in task_response.get("tasks", []) if task["taskArn"] == self.task_arn), None) + if task is None: + return 1 + status = task.get("lastStatus", None) + desired_status = task.get("desiredStatus", None) + stop_code = task.get("stopCode", None) + stopped_reason = task.get("stoppedReason", None) + stopped_at = task.get("stoppedAt", None) + + if status == "RUNNING" and desired_status == "RUNNING": + # TODO: Also check starting/provisioning statuses + return None + + self.log.warning(f"============\n\n{status=}\n\n{desired_status=}\n\n{stop_code=}\n\n{stopped_reason=}\n\n{stopped_at=}") + + try: + return int(stop_code) + except ValueError: + return 1 + + + + diff --git a/src/beakerhub/services/spawner/base.py b/src/beakerhub/services/spawner/base.py new file mode 100644 index 0000000..624e9da --- /dev/null +++ b/src/beakerhub/services/spawner/base.py @@ -0,0 +1,148 @@ +from typing import Any, cast +from uuid import uuid4 + +from jupyterhub.spawner import Spawner +from traitlets import Dict, Unicode, default, validate + +from beakerhub.auth.user import BeakerhubUser +from beakerhub.services.secrets import VAULT_ENV_VAR_LIST_KEY + +class BeakerSpawner(Spawner): + + default_beaker_context = Unicode( + default_value="default", + help="Slug of context the Beaker kernel should be started with." + ).tag(config=True) + + beaker_context = Unicode() + context_config = Dict() + node_env = Dict() + policy_override_env_key_suffix = Unicode( + default_value="_secret_policy", + help="", + config=True, + ) + node_policy_overrides = Dict( + help="Secret-handling policy overrides for this launch, keyed by environment " + "variable name, e.g. {'SHARED_API_KEY': {'subkernel_environment_policy': 'allow'}}. " + "Resolved from the vault at spawn time and passed to the node as JSON." + ) + + + 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 and self.orm_spawner.name != value: + self.orm_spawner.name = value + self.orm_spawner.save() + return value + + @property + def session_id(self) -> str: + if not self.name: + self.name = str(uuid4()) + return self.name + + def get_env(self): + 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) + env.setdefault("BEAKER_UI_HIDE_CONTEXT_SELECTOR", "true") + env.update(self.node_env) + # Tell the node which variables came from the vault. It discovers secrets by name + # heuristic, which misses anything innocuous-looking, and an unrecognized vault + # secret gets no policies applied *and* no default protection either. + if self.node_env: + env.setdefault(VAULT_ENV_VAR_LIST_KEY, ",".join(sorted(self.node_env))) + # Policy metadata is not itself sensitive, so it travels as a plain env var. + # Only set it when there is something to say, so the node keeps its own default + # of "no overrides" rather than parsing an empty object. + for env_name, policy_dict in self.node_policy_overrides.items(): + for policy_name, policy_value in policy_dict.items(): + policy_key = f"{env_name}_{policy_name}{self.policy_override_env_key_suffix}" + env.setdefault(policy_key, policy_value) + + return env + + + def start(self): + raise NotImplementedError() + + def stop(self, now=False): + raise NotImplementedError() + + def poll(self): + raise NotImplementedError() + + +class BeakerhubImageSpawner(BeakerSpawner): + + 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 diff --git a/src/beakerhub/spawner/kubernetes.py b/src/beakerhub/services/spawner/kubernetes_spawner.py similarity index 57% rename from src/beakerhub/spawner/kubernetes.py rename to src/beakerhub/services/spawner/kubernetes_spawner.py index 1bce24f..9b7b5db 100644 --- a/src/beakerhub/spawner/kubernetes.py +++ b/src/beakerhub/services/spawner/kubernetes_spawner.py @@ -6,47 +6,17 @@ from traitlets import default, validate, Unicode, Dict, List from traitlets.config import Application -from beakerhub.auth.user import BeakerhubUser -from beakerhub.spawner.base import BeakerSpawner -from beakerhub.services.secrets import VAULT_ENV_VAR_LIST_KEY +from beakerhub.services.spawner.base import BeakerSpawner, BeakerhubImageSpawner from beakerhub import orm -class BeakerKubeSpawner(KubeSpawner, BeakerSpawner): +class BeakerKubeSpawner(KubeSpawner, BeakerhubImageSpawner): - default_registry = Unicode().tag(config=True) - default_image = Unicode().tag(config=True) - default_tag = Unicode(default_value="latest").tag(config=True) + image = KubeSpawner.image - default_beaker_context = Unicode( - default_value="default", - help="Slug of context the Beaker kernel should be started with." - ).tag(config=True) - - beaker_context = Unicode() - context_config = Dict() - node_env = Dict() - policy_override_env_key_suffix = Unicode( - default_value="_secret_policy", - help="", - config=True, - ) - node_policy_overrides = Dict( - help="Secret-handling policy overrides for this launch, keyed by environment " - "variable name, e.g. {'SHARED_API_KEY': {'subkernel_environment_policy': 'allow'}}. " - "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("/") + 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): @@ -56,30 +26,6 @@ def _default_delete_stopped_pods(self): 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 - - @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 - @default("pod_name_template") def _default_pod_name_template(self): return "session-{user_server}" @@ -91,29 +37,6 @@ def node_type(self) -> str | None: return None return pod_name - def get_env(self): - user: BeakerhubUser = cast(BeakerhubUser, self.user) - env = super().get_env() - 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) - env.setdefault("BEAKER_UI_HIDE_CONTEXT_SELECTOR", "true") - env.update(self.node_env) - # Tell the node which variables came from the vault. It discovers secrets by name - # heuristic, which misses anything innocuous-looking, and an unrecognized vault - # secret gets no policies applied *and* no default protection either. - if self.node_env: - env.setdefault(VAULT_ENV_VAR_LIST_KEY, ",".join(sorted(self.node_env))) - # Policy metadata is not itself sensitive, so it travels as a plain env var. - # Only set it when there is something to say, so the node keeps its own default - # of "no overrides" rather than parsing an empty object. - for env_name, policy_dict in self.node_policy_overrides.items(): - for policy_name, policy_value in policy_dict.items(): - policy_key = f"{env_name}_{policy_name}{self.policy_override_env_key_suffix}" - env.setdefault(policy_key, policy_value) - - return env - @staticmethod def apply_user_options(spawner: "BeakerKubeSpawner", user_options: dict): node_record: orm.NodeImages | None = None 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..62be094 --- /dev/null +++ b/src/beakerhub/services/task/aws_ecs_task_runner.py @@ -0,0 +1,26 @@ +"""AWS ECS implementation skeleton for the task-runner service.""" + +from typing import TYPE_CHECKING + +from beakerhub.services.task.base import BaseTaskRunnerService + +if TYPE_CHECKING: + from beakerhub.orm import NodeImages + + +class AwsEcsTaskRunnerService(BaseTaskRunnerService): + """Run BeakerHub background tasks as AWS ECS tasks.""" + + def submit_image_import( + self, + node_image: "NodeImages", + callback_url: str, + callback_token: str, + ) -> str: + raise NotImplementedError("AWS ECS task submission is not implemented") + + def get_status(self, task_id: str) -> dict: + raise NotImplementedError("AWS ECS task status polling is not implemented") + + def delete(self, task_id: str) -> None: + raise NotImplementedError("AWS ECS task deletion is not implemented") diff --git a/src/beakerhub/services/task/base.py b/src/beakerhub/services/task/base.py new file mode 100644 index 0000000..eea37a0 --- /dev/null +++ b/src/beakerhub/services/task/base.py @@ -0,0 +1,29 @@ +"""Base contract for background task-runner services.""" + +from typing import TYPE_CHECKING + +from traitlets.config import LoggingConfigurable + +if TYPE_CHECKING: + from beakerhub.orm import NodeImages + + +class BaseTaskRunnerService(LoggingConfigurable): + """Submit and manage background workloads for BeakerHub tasks.""" + + def submit_image_import( + self, + node_image: "NodeImages", + callback_url: str, + callback_token: str, + ) -> str: + """Submit an image-import workload and return its external identifier.""" + raise NotImplementedError + + def get_status(self, task_id: str) -> dict: + """Return the current state and diagnostic message for a workload.""" + raise NotImplementedError + + def delete(self, task_id: str) -> None: + """Delete or cancel a workload.""" + raise NotImplementedError diff --git a/src/beakerhub/nodes/import_handlers.py b/src/beakerhub/services/task/handlers.py similarity index 63% rename from src/beakerhub/nodes/import_handlers.py rename to src/beakerhub/services/task/handlers.py index b05dbbd..c6aef76 100644 --- a/src/beakerhub/nodes/import_handlers.py +++ b/src/beakerhub/services/task/handlers.py @@ -1,120 +1,25 @@ -""" -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) -""" +"""API handlers for background tasks.""" import json import logging from datetime import datetime, timezone -from secrets import token_urlsafe -from typing import Any, TYPE_CHECKING +from typing import Any 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 +from beakerhub.tasks.image_import.task import launch_import_task +from beakerhub.utils import ingest_interchange_dump 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.""" + """Create a task 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}") @@ -159,10 +64,11 @@ async def get(self, image_id: str): 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 the task is still running, poll the configured task runner. if task.status == "running" and task.job_name: try: - job_status = k8s_tasks.get_job_status(task.job_name) + app = self.settings["app"] + job_status = app.task_runner.get_status(task.job_name) if job_status["status"] == "failed": task.status = "failed" task.error = job_status["message"] @@ -196,7 +102,7 @@ class TaskCallbackHandler(APIHandler): """ def check_xsrf_cookie(self): - # Internal endpoint called from K8s pod, no XSRF cookie + # Internal endpoint called by a task workload, with no XSRF cookie. return def get_current_user(self): @@ -264,8 +170,10 @@ async def post(self, callback_token: str): self.write(json.dumps({"status": "ok", "stats": combined_stats})) -import_handlers = [ +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), ] + +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..2722856 --- /dev/null +++ b/src/beakerhub/services/task/kubernetes_task_runner.py @@ -0,0 +1,264 @@ +"""Kubernetes implementation of the task-runner service.""" + +from typing import Any, TYPE_CHECKING +from uuid import uuid4 + +from kubernetes import client as k8s_client +from kubernetes import config as k8s_config +from traitlets import Dict, Integer, List, Unicode + +from beakerhub.services.task.base import BaseTaskRunnerService + +if TYPE_CHECKING: + from beakerhub.orm import NodeImages + + +class KubernetesTaskRunnerService(BaseTaskRunnerService): + """Run BeakerHub background tasks as Kubernetes Jobs.""" + + reporter_image = Unicode( + "beakerhub/task-reporter:latest", + config=True, + help="Full image reference for the task reporter container.", + ) + reporter_pull_policy = Unicode( + "Always", + config=True, + help="Image pull policy for the task reporter container.", + ) + reporter_resources = Dict( + config=True, + help="Resource requests and limits for the task reporter container.", + ) + node_image_resources = Dict( + config=True, + help="Resource requests and limits for the node-image init 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( + "beakerhub", + config=True, + help="Kubernetes namespace for task Jobs.", + ) + + @staticmethod + def _get_clients() -> tuple[k8s_client.BatchV1Api, k8s_client.CoreV1Api]: + try: + k8s_config.load_incluster_config() + except k8s_config.ConfigException: + k8s_config.load_kube_config() + return k8s_client.BatchV1Api(), k8s_client.CoreV1Api() + + def submit_image_import( + self, + node_image: "NodeImages", + callback_url: str, + callback_token: str, + ) -> str: + """Create a Job that extracts and reports context data from an image.""" + batch_api, _ = self._get_clients() + job_name = f"node-import-{node_image.slug}-{uuid4().hex[:8]}" + output_volume_name = "task-output" + stdout_path = "/output/stdout" + stderr_path = "/output/stderr" + + init_container = k8s_client.V1Container( + name="context-dump", + image=node_image.default_img_string, + 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=self._build_resource_requirements(self.node_image_resources), + ) + reporter_container = k8s_client.V1Container( + name="reporter", + image=self.reporter_image, + image_pull_policy=self.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=self._build_resource_requirements(self.reporter_resources), + ) + 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=self.node_selector or None, + tolerations=( + [k8s_client.V1Toleration(**item) for item in self.tolerations] + if self.tolerations + else None + ), + ) + job = k8s_client.V1Job( + api_version="batch/v1", + kind="Job", + metadata=k8s_client.V1ObjectMeta( + name=job_name, + namespace=self.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=self.backoff_limit, + active_deadline_seconds=self.active_deadline_seconds, + ttl_seconds_after_finished=self.ttl_seconds_after_finished, + ), + ) + + batch_api.create_namespaced_job(namespace=self.namespace, body=job) + self.log.info( + "Created import job %s for node image %s", job_name, node_image.slug + ) + return job_name + + def get_status(self, task_id: str) -> dict: + """Query the status of a Kubernetes Job.""" + batch_api, core_api = self._get_clients() + try: + job = batch_api.read_namespaced_job( + name=task_id, + namespace=self.namespace, + ) + except k8s_client.ApiException as error: + if error.status == 404: + return { + "status": "failed", + "message": f"Job {task_id} 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: + return { + "status": "failed", + "message": self._get_failure_message(core_api, task_id), + } + if status.active and status.active > 0: + return {"status": "running", "message": "Job is running"} + return {"status": "pending", "message": "Job is pending"} + + def _get_failure_message( + self, + core_api: k8s_client.CoreV1Api, + job_name: str, + ) -> str: + try: + pods = core_api.list_namespaced_pod( + namespace=self.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] + for status in pod.status.init_container_statuses or []: + message = self._container_failure_message(status, "Init container") + if message: + return message + for status in pod.status.container_statuses or []: + message = self._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() + + @staticmethod + def _container_failure_message(status, label: str) -> str | None: + 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 + + def delete(self, task_id: str) -> None: + """Delete a Job and its associated Pods.""" + batch_api, _ = self._get_clients() + try: + batch_api.delete_namespaced_job( + name=task_id, + namespace=self.namespace, + body=k8s_client.V1DeleteOptions(propagation_policy="Background"), + ) + except k8s_client.ApiException as error: + if error.status != 404: + raise + + @staticmethod + def _build_resource_requirements( + resources: dict, + ) -> k8s_client.V1ResourceRequirements | None: + if not resources: + return None + return k8s_client.V1ResourceRequirements( + requests=resources.get("requests"), + limits=resources.get("limits"), + ) 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/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..3fef4b2 --- /dev/null +++ b/src/beakerhub/tasks/image_import/task.py @@ -0,0 +1,78 @@ +"""Orchestration for importing context metadata from a node image.""" + +import logging +from datetime import datetime, timezone +from secrets import token_urlsafe +from typing import TYPE_CHECKING + +from sqlalchemy.orm import Session + +from beakerhub.orm import NodeImages, NodeImageTask + +if TYPE_CHECKING: + from beakerhub.app import BeakerHub + + +log = logging.getLogger(__name__) + + +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})" + ) + + 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() + + hub_url = getattr(app, "hub_connect_url", "http://localhost:8888") + callback_url = ( + f"{hub_url.rstrip('/')}/api/beakerhub/internal/task-callback/" + f"{callback_token}" + ) + try: + external_task_id = app.task_runner.submit_image_import( + node_image=node_image, + callback_url=callback_url, + callback_token=callback_token, + ) + 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 = external_task_id + task.status = "running" + task.updated_at = datetime.now(timezone.utc) + db.commit() + log.info( + "Launched import task %s for node image %s", + external_task_id, + node_image.slug, + ) + return task 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/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 99% rename from tests/unit/spawner/test_kubernetes.py rename to tests/unit/services/spawner/test_kubernetes.py index 0782009..8ce56ef 100644 --- a/tests/unit/spawner/test_kubernetes.py +++ b/tests/unit/services/spawner/test_kubernetes.py @@ -1,13 +1,13 @@ # 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 class TestBeakerKubeSpawner: diff --git a/tests/unit/services/test_service_wiring.py b/tests/unit/services/test_service_wiring.py new file mode 100644 index 0000000..5ccd471 --- /dev/null +++ b/tests/unit/services/test_service_wiring.py @@ -0,0 +1,64 @@ +"""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.services.dashboard.base import BaseDashboardService +from beakerhub.services.task.base import BaseTaskRunnerService +from beakerhub.tasks.image_import.task import 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_task_runner_receives_config_loaded_after_its_creation(): + app = BeakerHub() + runner = app.task_runner + config = Config() + config.KubernetesTaskRunnerService.reporter_image = "registry.example/reporter:v1" + + app.update_config(config) + + assert runner.reporter_image == "registry.example/reporter:v1" + + +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_image_import.return_value = "external-task-id" + + task = launch_import_task(db, app, node_image) + + app.task_runner.submit_image_import.assert_called_once() + assert task.job_name == "external-task-id" + assert task.status == "running" + assert db.commit.call_count == 2 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 From 47b2e87aca768f744a929c26ac3d53f5207965e6 Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Tue, 25 Aug 2026 16:11:09 -0600 Subject: [PATCH 2/8] Updates for end-to-end ecs spawner initial buildout --- docs/DEVELOPMENT.md | 6 +- src/beakerhub/services/spawner/__init__.py | 3 +- .../services/spawner/aws_ecs_spawner.py | 210 +++++++++++------- tests/unit/services/spawner/test_aws_ecs.py | 150 +++++++++++++ 4 files changed, 284 insertions(+), 85 deletions(-) create mode 100644 tests/unit/services/spawner/test_aws_ecs.py 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/src/beakerhub/services/spawner/__init__.py b/src/beakerhub/services/spawner/__init__.py index 1146626..80d6fe7 100644 --- a/src/beakerhub/services/spawner/__init__.py +++ b/src/beakerhub/services/spawner/__init__.py @@ -1,5 +1,6 @@ """Spawner services for BeakerHub session runtimes.""" +from .aws_ecs_spawner import BeakerAwsECSSpawner from .base import BeakerSpawner, BeakerhubImageSpawner -__all__ = ["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 index 7cab0ae..9067706 100644 --- a/src/beakerhub/services/spawner/aws_ecs_spawner.py +++ b/src/beakerhub/services/spawner/aws_ecs_spawner.py @@ -1,107 +1,153 @@ -import json -from typing import cast, Optional +"""ECS-backed JupyterHub spawner. + +This module deliberately contains only ECS control-plane behavior. Moto can +exercise that behavior in unit tests; a Docker-backed ECS implementation such +as LocalStack is needed to verify that a task image actually starts. +""" + +from typing import Any import boto3 -from jupyterhub.spawner import Spawner -from kubespawner.spawner import KubeSpawner -from traitlets import default, validate, Unicode, Dict, List, Bool -from traitlets.config import Application +from traitlets import Bool, Dict, List, Unicode -from beakerhub.auth.user import BeakerhubUser -from beakerhub.services.spawner.base import BeakerSpawner, BeakerhubImageSpawner -from beakerhub.services.secrets import VAULT_ENV_VAR_LIST_KEY -from beakerhub import orm +from beakerhub.services.spawner.base import BeakerhubImageSpawner class BeakerAwsECSSpawner(BeakerhubImageSpawner): - cluster_name: str = Unicode( + """Launch a Beaker session as one ECS task. + + Networking/proxy registration is intentionally not implemented here yet. + ``start`` therefore launches and records the ECS task but does not claim + that the Jupyter server is reachable. + """ + + cluster_name = Unicode("default", config=True, help="ECS cluster name or ARN.") + task_definition = Unicode( config=True, + help="ECS task-definition family, revision, or ARN for notebook tasks.", ) - container_type: str = Unicode( + container_name = Unicode( config=True, + help="Name of the notebook container in the ECS task definition.", ) - container_tags: dict[str, str] = Dict( - Unicode, - config=True, + container_tags = Dict( + Unicode(), Unicode(), default_value={}, config=True, + help="Tags applied to every ECS notebook task.", ) - fargate: bool = Bool( - False, - config=True, + subnets = List( + Unicode(), default_value=[], config=True, + help="Subnets used by the task's awsvpc network configuration.", ) + security_groups = List(Unicode(), default_value=[], config=True) + assign_public_ip = Bool(False, config=True) + fargate = Bool(False, config=True) - boto: boto3.client - # session_key: str - task_arn = Optional[str] - - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.boto = boto3.client("ecs") - self.task_arn = None + self.task_arn: str | None = None @property - def volume_configs(self) -> list[dict]: + def volume_configs(self) -> list[dict[str, Any]]: + """Return ECS managed-volume configuration overrides, if configured.""" return [] + def _task_tags(self) -> list[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 [{"key": key, "value": value} for key, value in tags.items()] - def start(self): - overrides = {} - tasks = self.boto.run_task( - cluster=self.cluster_name, - count=1, - networkConfiguration={ - "assignPublicIp": "ENABLED", + def _network_configuration(self) -> dict[str, Any] | None: + if not self.subnets: + if self.fargate: + raise ValueError("BeakerAwsECSSpawner.subnets is required for Fargate tasks") + return None + return { + "awsvpcConfiguration": { + "subnets": self.subnets, + "securityGroups": self.security_groups, + "assignPublicIp": "ENABLED" if self.assign_public_ip else "DISABLED", + } + } + + def _run_task_request(self) -> dict[str, Any]: + if not self.task_definition: + raise ValueError("BeakerAwsECSSpawner.task_definition must be configured") + if not self.container_name: + raise ValueError("BeakerAwsECSSpawner.container_name must be configured") + + request: dict[str, Any] = { + "cluster": self.cluster_name, + "taskDefinition": self.task_definition, + "count": 1, + "overrides": { + "containerOverrides": [{ + "name": self.container_name, + "environment": [ + {"name": name, "value": str(value)} + for name, value in self.get_env().items() + ], + }], }, - overrides=overrides, - tags=(self.container_tags + [{"key": "beaker-session", "value": self.session_id}]), - volumeConfigurations=self.volume_configs - ) - self.tasks = tasks - - def stop(self, now=False): - self.boto.stop_task( - cluster=self.cluster_name, - task=self.task_arn, - reason="Stopped by BeakerHub server" - ) - - - def poll(self): - """ - Check if the pod is still running. - - Uses the same interface as subprocess.Popen.poll(): if the pod is - still running, returns None. If the pod has exited, return the - exit code if we can determine it, or 1 if it has exited but we - don't know how. These are the return values JupyterHub expects. - - Note that a clean exit will have an exit code of zero, so it is - necessary to check that the returned value is None, rather than - just Falsy, to determine that the pod is still running. - """ - task_response = self.boto.describe_tasks( - cluster=self.cluster_name, - tasks=[self.task_arn], + "tags": self._task_tags(), + } + if self.fargate: + request["launchType"] = "FARGATE" + network_configuration = self._network_configuration() + if network_configuration: + request["networkConfiguration"] = network_configuration + if self.volume_configs: + request["volumeConfigurations"] = self.volume_configs + return request + + def start(self) -> None: + """Request one ECS task and retain its ARN for later lifecycle calls.""" + response = self.boto.run_task(**self._run_task_request()) + tasks = response.get("tasks", []) + if not tasks: + failures = response.get("failures", []) + raise RuntimeError(f"ECS did not start a task: {failures!r}") + self.task_arn = tasks[0]["taskArn"] + + def stop(self, now: bool = False) -> None: + """Stop the task, if this spawner has launched one.""" + if self.task_arn: + self.boto.stop_task( + cluster=self.cluster_name, + task=self.task_arn, + reason="Stopped by BeakerHub server", + ) + + def poll(self) -> int | None: + """Return ``None`` while ECS is provisioning/running, otherwise an exit code.""" + if not self.task_arn: + return 0 + response = self.boto.describe_tasks(cluster=self.cluster_name, tasks=[self.task_arn]) + task = next( + (item for item in response.get("tasks", []) if item.get("taskArn") == self.task_arn), + None, ) - task = next((task for task in task_response.get("tasks", []) if task["taskArn"] == self.task_arn), None) if task is None: return 1 - status = task.get("lastStatus", None) - desired_status = task.get("desiredStatus", None) - stop_code = task.get("stopCode", None) - stopped_reason = task.get("stoppedReason", None) - stopped_at = task.get("stoppedAt", None) - - if status == "RUNNING" and desired_status == "RUNNING": - # TODO: Also check starting/provisioning statuses + if task.get("lastStatus") in {"PROVISIONING", "PENDING", "ACTIVATING", "RUNNING"}: return None - - self.log.warning(f"============\n\n{status=}\n\n{desired_status=}\n\n{stop_code=}\n\n{stopped_reason=}\n\n{stopped_at=}") - - try: - return int(stop_code) - except ValueError: - return 1 - - - - + for container in task.get("containers", []): + exit_code = container.get("exitCode") + if exit_code is not None: + return int(exit_code) + return 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 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..72da187 --- /dev/null +++ b/tests/unit/services/spawner/test_aws_ecs.py @@ -0,0 +1,150 @@ +"""Unit tests for the ECS control-plane spawner scaffold. + +These use botocore's request-validating Stubber rather than Docker: they are +fast tests for the ECS contract. A future LocalStack suite should cover actual +container execution separately. +""" + +from types import SimpleNamespace + +import boto3 +import pytest +from botocore.stub import Stubber + +from beakerhub.services.spawner.aws_ecs_spawner import BeakerAwsECSSpawner + + +TASK_ARN = "arn:aws:ecs:us-east-1:123456789012:task/test/task-123" + + +def configured_spawner(**overrides): + """A minimal object sufficient to exercise unbound ECS-spawner methods.""" + values = { + "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, + "fargate": True, + "volume_configs": [], + "get_env": lambda: {"BEAKERHUB_USER": "ada", "NUMBER": 2}, + } + values.update(overrides) + spawner = SimpleNamespace(**values) + spawner._task_tags = lambda: BeakerAwsECSSpawner._task_tags(spawner) + spawner._network_configuration = lambda: BeakerAwsECSSpawner._network_configuration(spawner) + return spawner + + +def test_run_task_request_contains_environment_tags_and_fargate_networking(): + spawner = configured_spawner() + + request = BeakerAwsECSSpawner._run_task_request(spawner) + + assert request == { + "cluster": "notebooks", + "taskDefinition": "beaker-notebook:3", + "count": 1, + "launchType": "FARGATE", + "networkConfiguration": { + "awsvpcConfiguration": { + "subnets": ["subnet-123"], + "securityGroups": ["sg-123"], + "assignPublicIp": "ENABLED", + } + }, + "overrides": { + "containerOverrides": [{ + "name": "notebook", + "environment": [ + {"name": "BEAKERHUB_USER", "value": "ada"}, + {"name": "NUMBER", "value": "2"}, + ], + }] + }, + "tags": [ + {"key": "environment", "value": "test"}, + {"key": "beaker-session", "value": "session-123"}, + ], + } + + +def test_session_tag_cannot_be_overridden_by_configuration(): + spawner = configured_spawner(container_tags={"beaker-session": "incorrect"}) + + assert BeakerAwsECSSpawner._task_tags(spawner) == [ + {"key": "beaker-session", "value": "session-123"} + ] + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"task_definition": ""}, "task_definition"), + ({"container_name": ""}, "container_name"), + ({"subnets": []}, "subnets"), + ], +) +def test_run_task_request_rejects_incomplete_fargate_configuration(overrides, message): + spawner = configured_spawner(**overrides) + + with pytest.raises(ValueError, match=message): + BeakerAwsECSSpawner._run_task_request(spawner) + + +def test_start_stop_and_poll_use_ecs_control_plane(): + client = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + request = {"cluster": "notebooks", "taskDefinition": "beaker-notebook:3", "count": 1} + spawner = configured_spawner(boto=client, task_arn=None) + spawner._run_task_request = lambda: request + + with Stubber(client) as stubber: + stubber.add_response("run_task", {"tasks": [{"taskArn": TASK_ARN}]}, request) + BeakerAwsECSSpawner.start(spawner) + assert spawner.task_arn == TASK_ARN + + stubber.add_response( + "describe_tasks", + {"tasks": [{"taskArn": TASK_ARN, "lastStatus": "RUNNING"}]}, + {"cluster": "notebooks", "tasks": [TASK_ARN]}, + ) + assert BeakerAwsECSSpawner.poll(spawner) is None + + stubber.add_response( + "stop_task", + {"task": {"taskArn": TASK_ARN, "lastStatus": "STOPPED"}}, + { + "cluster": "notebooks", + "task": TASK_ARN, + "reason": "Stopped by BeakerHub server", + }, + ) + BeakerAwsECSSpawner.stop(spawner) + + +def test_poll_returns_container_exit_code_for_stopped_task(): + client = boto3.client( + "ecs", region_name="us-east-1", aws_access_key_id="test", aws_secret_access_key="test" + ) + spawner = configured_spawner(boto=client, task_arn=TASK_ARN) + with Stubber(client) as stubber: + stubber.add_response( + "describe_tasks", + { + "tasks": [{ + "taskArn": TASK_ARN, + "lastStatus": "STOPPED", + "containers": [{"name": "notebook", "exitCode": 23}], + }] + }, + {"cluster": "notebooks", "tasks": [TASK_ARN]}, + ) + assert BeakerAwsECSSpawner.poll(spawner) == 23 From 4775c219ed205e57f3037bd660ac9cbabe8947dd Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Thu, 27 Aug 2026 17:16:00 -0600 Subject: [PATCH 3/8] Complete refactor of task running --- .../beakerhub/templates/beakerhub-config.yaml | 3 - src/beakerhub/orm.py | 22 ++- .../services/task/aws_ecs_task_runner.py | 21 +-- src/beakerhub/services/task/base.py | 95 +++++++++-- src/beakerhub/services/task/handlers.py | 134 ++++----------- .../services/task/kubernetes_task_runner.py | 156 ++++++++---------- src/beakerhub/tasks/base.py | 54 ++++++ src/beakerhub/tasks/image_import/task.py | 64 +++++-- tests/unit/services/test_service_wiring.py | 18 +- tests/unit/services/test_task_base.py | 57 +++++++ 10 files changed, 386 insertions(+), 238 deletions(-) create mode 100644 src/beakerhub/tasks/base.py create mode 100644 tests/unit/services/test_task_base.py diff --git a/helm/beakerhub/templates/beakerhub-config.yaml b/helm/beakerhub/templates/beakerhub-config.yaml index 4f73d9b..b0d8746 100644 --- a/helm/beakerhub/templates/beakerhub-config.yaml +++ b/helm/beakerhub/templates/beakerhub-config.yaml @@ -192,9 +192,6 @@ c.Authenticator.admin_users = {{ .Values.auth.adminUsers | toJson }} #------------------------------------------------------------------------------ c.KubernetesTaskRunnerService.namespace = {{ .Values.namespace | quote }} -c.KubernetesTaskRunnerService.reporter_image = {{ printf "%s%s:%s" (default "" .Values.defaultRegistry) .Values.tasks.reporter.image.repository .Values.tasks.reporter.image.tag | quote }} -c.KubernetesTaskRunnerService.reporter_pull_policy = {{ .Values.tasks.reporter.image.pullPolicy | quote }} -c.KubernetesTaskRunnerService.reporter_resources = {{ .Values.tasks.reporter.resources | toJson }} 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 }} 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/services/task/aws_ecs_task_runner.py b/src/beakerhub/services/task/aws_ecs_task_runner.py index 62be094..f774f6d 100644 --- a/src/beakerhub/services/task/aws_ecs_task_runner.py +++ b/src/beakerhub/services/task/aws_ecs_task_runner.py @@ -1,25 +1,20 @@ """AWS ECS implementation skeleton for the task-runner service.""" -from typing import TYPE_CHECKING - -from beakerhub.services.task.base import BaseTaskRunnerService - -if TYPE_CHECKING: - from beakerhub.orm import NodeImages +from beakerhub.services.task.base import ( + BaseTaskRunnerService, + RunningTask, + TaskStatus, +) +from beakerhub.tasks.base import BaseTaskDefinition class AwsEcsTaskRunnerService(BaseTaskRunnerService): """Run BeakerHub background tasks as AWS ECS tasks.""" - def submit_image_import( - self, - node_image: "NodeImages", - callback_url: str, - callback_token: str, - ) -> str: + def submit(self, task: BaseTaskDefinition) -> RunningTask: raise NotImplementedError("AWS ECS task submission is not implemented") - def get_status(self, task_id: str) -> dict: + def get_status(self, task_id: str) -> TaskStatus: raise NotImplementedError("AWS ECS task status polling is not implemented") def delete(self, task_id: str) -> None: diff --git a/src/beakerhub/services/task/base.py b/src/beakerhub/services/task/base.py index eea37a0..bcad92f 100644 --- a/src/beakerhub/services/task/base.py +++ b/src/beakerhub/services/task/base.py @@ -1,29 +1,96 @@ -"""Base contract for background task-runner services.""" +"""Base contracts for background task-runner services.""" -from typing import TYPE_CHECKING +import asyncio +from dataclasses import dataclass +from time import monotonic +from typing import TYPE_CHECKING, Literal from traitlets.config import LoggingConfigurable if TYPE_CHECKING: - from beakerhub.orm import NodeImages + from beakerhub.tasks.base import BaseTaskDefinition + + +TaskState = Literal["pending", "running", "completed", "failed"] + + +@dataclass(frozen=True) +class TaskStatus: + """The runner's current view of a submitted task.""" + + state: TaskState + message: str | None = None + + @property + def done(self) -> bool: + """Return true when the task has reached a terminal state.""" + return self.state in {"completed", "failed"} + + +@dataclass(frozen=True) +class TaskOutput: + """Captured standard streams for a completed task.""" + + stdout: str + stderr: str + + +@dataclass(frozen=True) +class RunningTask: + """An in-memory handle for a task submitted to a runner.""" + + external_id: str + task_definition: "BaseTaskDefinition" + runner: "BaseTaskRunnerService" + + @property + def status(self) -> TaskStatus: + """Return the current status from the runner.""" + return self.runner.get_status(self.external_id) + + @property + def done(self) -> bool: + """Return true when the runner reports a terminal state.""" + return self.status.done + + @property + def output(self) -> TaskOutput | None: + """Return captured output after the runner reports task completion.""" + return self.runner.get_output(self.external_id) + + async def await_completion(self, timeout: float | None = 600) -> TaskStatus: + """Wait for terminal status, or raise ``TimeoutError``.""" + 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"Task {self.external_id!r} did not complete within {timeout} seconds" + ) + await asyncio.sleep(0.2) class BaseTaskRunnerService(LoggingConfigurable): """Submit and manage background workloads for BeakerHub tasks.""" - def submit_image_import( - self, - node_image: "NodeImages", - callback_url: str, - callback_token: str, - ) -> str: - """Submit an image-import workload and return its external identifier.""" + def submit(self, task: "BaseTaskDefinition") -> RunningTask: + """Submit a task and return its in-memory runtime handle.""" + raise NotImplementedError + + def get_status(self, external_id: str) -> TaskStatus: + """Return the current state and diagnostics for a submitted task.""" + raise NotImplementedError + + def get_output(self, external_id: str) -> TaskOutput | None: + """Return output for a terminal task, if the backend retains it.""" raise NotImplementedError - def get_status(self, task_id: str) -> dict: - """Return the current state and diagnostic message for a workload.""" + def delete(self, external_id: str) -> None: + """Remove a known submitted workload during normal task completion.""" raise NotImplementedError - def delete(self, task_id: str) -> None: - """Delete or cancel a workload.""" + 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 index c6aef76..6b35d7e 100644 --- a/src/beakerhub/services/task/handlers.py +++ b/src/beakerhub/services/task/handlers.py @@ -1,15 +1,17 @@ """API handlers for background tasks.""" + import json import logging from datetime import datetime, timezone from typing import Any -from tornado import web 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 launch_import_task -from beakerhub.utils import ingest_interchange_dump +from beakerhub.tasks.image_import.task import ingest_import_output, launch_import_task + log = logging.getLogger(__name__) @@ -17,20 +19,18 @@ class NodeImageImportHandler(APIHandler): """Trigger a context import for a node image.""" - @needs_scope('admin:users') + @needs_scope("admin:users") async def post(self, image_id: str): - """Create a task 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)) + 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({ @@ -41,19 +41,17 @@ async def post(self, image_id: str): class NodeImageImportStatusHandler(APIHandler): - """Poll the status of an import task.""" + """Poll and collect the result of a node-image import task.""" def compute_etag(self) -> None: return None - @needs_scope('admin:users') + @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", @@ -64,18 +62,32 @@ async def get(self, image_id: str): self.write(json.dumps({"status": "none", "message": "No import has been run"})) return - # If the task is still running, poll the configured task runner. if task.status == "running" and task.job_name: try: - app = self.settings["app"] - job_status = app.task_runner.get_status(task.job_name) - if job_status["status"] == "failed": + runner = self.settings["app"].task_runner + status = runner.get_status(task.job_name) + if status.state == "completed": + output = runner.get_output(task.job_name) + if output is None: + raise RuntimeError("Task completed but no output is available") + task.result = ingest_import_output(self.db, node, 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) + messages = [status.message] + if output and output.stderr: + messages.append(output.stderr) task.status = "failed" - task.error = job_status["message"] + task.error = ". ".join(message for message in messages if 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}") + 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, @@ -93,87 +105,9 @@ async def get(self, image_id: str): 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 by a task workload, with 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})) - - 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), ] api_handlers = handlers diff --git a/src/beakerhub/services/task/kubernetes_task_runner.py b/src/beakerhub/services/task/kubernetes_task_runner.py index 2722856..557284f 100644 --- a/src/beakerhub/services/task/kubernetes_task_runner.py +++ b/src/beakerhub/services/task/kubernetes_task_runner.py @@ -1,35 +1,24 @@ """Kubernetes implementation of the task-runner service.""" -from typing import Any, TYPE_CHECKING +from typing import Any from uuid import uuid4 from kubernetes import client as k8s_client from kubernetes import config as k8s_config from traitlets import Dict, Integer, List, Unicode -from beakerhub.services.task.base import BaseTaskRunnerService - -if TYPE_CHECKING: - from beakerhub.orm import NodeImages +from beakerhub.services.task.base import ( + BaseTaskRunnerService, + RunningTask, + TaskOutput, + TaskStatus, +) +from beakerhub.tasks.base import BaseImageTask, BaseTaskDefinition class KubernetesTaskRunnerService(BaseTaskRunnerService): """Run BeakerHub background tasks as Kubernetes Jobs.""" - reporter_image = Unicode( - "beakerhub/task-reporter:latest", - config=True, - help="Full image reference for the task reporter container.", - ) - reporter_pull_policy = Unicode( - "Always", - config=True, - help="Image pull policy for the task reporter container.", - ) - reporter_resources = Dict( - config=True, - help="Resource requests and limits for the task reporter container.", - ) node_image_resources = Dict( config=True, help="Resource requests and limits for the node-image init container.", @@ -65,62 +54,41 @@ def _get_clients() -> tuple[k8s_client.BatchV1Api, k8s_client.CoreV1Api]: k8s_config.load_kube_config() return k8s_client.BatchV1Api(), k8s_client.CoreV1Api() - def submit_image_import( - self, - node_image: "NodeImages", - callback_url: str, - callback_token: str, - ) -> str: - """Create a Job that extracts and reports context data from an image.""" - batch_api, _ = self._get_clients() - job_name = f"node-import-{node_image.slug}-{uuid4().hex[:8]}" - output_volume_name = "task-output" - stdout_path = "/output/stdout" - stderr_path = "/output/stderr" - - init_container = k8s_client.V1Container( - name="context-dump", - image=node_image.default_img_string, - 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=self._build_resource_requirements(self.node_image_resources), + def submit(self, task: BaseTaskDefinition) -> RunningTask: + """Submit a supported task as a Kubernetes Job.""" + if not isinstance(task, BaseImageTask): + raise ValueError( + f"KubernetesTaskRunnerService only supports image tasks, not " + f"{task.task_type!r}" + ) + external_id = self._submit_image_task(task) + return RunningTask( + external_id=external_id, + task_definition=task, + runner=self, ) - reporter_container = k8s_client.V1Container( - name="reporter", - image=self.reporter_image, - image_pull_policy=self.reporter_pull_policy, + + def _submit_image_task(self, task: BaseImageTask) -> str: + """Create a Job for a runtime-neutral image-task definition.""" + batch_api, _ = self._get_clients() + task_name = task.task_type.replace("_", "-") + job_name = f"beaker-task-{task_name}-{uuid4().hex[:8]}" + container = k8s_client.V1Container( + name="task", + image=task.image, + command=list(task.entrypoint) or None, + args=list(task.command) or None, + working_dir=task.working_directory, 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=self._build_resource_requirements(self.reporter_resources), + k8s_client.V1EnvVar(name=name, value=value) + for name, value in task.environment.items() + ] or None, + resources=self._build_resource_requirements( + dict(task.resources) or self.node_image_resources + ), ) 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(), - ) - ], + containers=[container], restart_policy="Never", node_selector=self.node_selector or None, tolerations=( @@ -137,9 +105,8 @@ def submit_image_import( namespace=self.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, + "app.kubernetes.io/component": "task", + "beakerhub/task-type": task.task_type, }, ), spec=k8s_client.V1JobSpec( @@ -147,7 +114,7 @@ def submit_image_import( metadata=k8s_client.V1ObjectMeta( labels={ "app.kubernetes.io/name": "beakerhub", - "app.kubernetes.io/component": "node-image-task", + "app.kubernetes.io/component": "task", } ), spec=pod_spec, @@ -160,11 +127,11 @@ def submit_image_import( batch_api.create_namespaced_job(namespace=self.namespace, body=job) self.log.info( - "Created import job %s for node image %s", job_name, node_image.slug + "Created task job %s for task type %s", job_name, task.task_type ) return job_name - def get_status(self, task_id: str) -> dict: + def get_status(self, task_id: str) -> TaskStatus: """Query the status of a Kubernetes Job.""" batch_api, core_api = self._get_clients() try: @@ -174,23 +141,36 @@ def get_status(self, task_id: str) -> dict: ) except k8s_client.ApiException as error: if error.status == 404: - return { - "status": "failed", - "message": f"Job {task_id} not found", - } + return TaskStatus("failed", f"Job {task_id} not found") raise status = job.status if status.succeeded and status.succeeded > 0: - return {"status": "completed", "message": "Job completed successfully"} + return TaskStatus("completed", "Job completed successfully") if status.failed and status.failed > 0: - return { - "status": "failed", - "message": self._get_failure_message(core_api, task_id), - } + return TaskStatus("failed", self._get_failure_message(core_api, task_id)) if status.active and status.active > 0: - return {"status": "running", "message": "Job is running"} - return {"status": "pending", "message": "Job is pending"} + return TaskStatus("running", "Job is running") + return TaskStatus("pending", "Job is pending") + + def get_output(self, task_id: str) -> TaskOutput | None: + """Return the task container log after its Job reaches a terminal state.""" + _, core_api = self._get_clients() + pods = core_api.list_namespaced_pod( + namespace=self.namespace, + label_selector=f"job-name={task_id}", + ) + if not pods.items: + return None + response = core_api.read_namespaced_pod_log( + name=pods.items[0].metadata.name, + namespace=self.namespace, + container="task", + _preload_content=False, + ) + logs = response.data.decode("utf-8") + # Kubernetes exposes a combined container log stream through this API. + return TaskOutput(stdout=logs or "", stderr="") def _get_failure_message( self, 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/task.py b/src/beakerhub/tasks/image_import/task.py index 3fef4b2..1964e9a 100644 --- a/src/beakerhub/tasks/image_import/task.py +++ b/src/beakerhub/tasks/image_import/task.py @@ -1,13 +1,17 @@ """Orchestration for importing context metadata from a node image.""" +import json import logging +from dataclasses import dataclass from datetime import datetime, timezone -from secrets import token_urlsafe from typing import TYPE_CHECKING from sqlalchemy.orm import Session from beakerhub.orm import NodeImages, NodeImageTask +from beakerhub.services.task.base import TaskOutput +from beakerhub.tasks.base import BaseImageTask +from beakerhub.utils import ingest_interchange_dump if TYPE_CHECKING: from beakerhub.app import BeakerHub @@ -16,6 +20,17 @@ log = logging.getLogger(__name__) +@dataclass(frozen=True, kw_only=True) +class ImageImportTask(BaseImageTask): + """Definition of a task that imports context metadata from a node image.""" + + node_image: NodeImages + + @property + def task_type(self) -> str: + return "context_import" + + def launch_import_task( db: Session, app: "BeakerHub", @@ -37,26 +52,22 @@ def launch_import_task( f"(job: {existing_task.job_name})" ) - 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() - hub_url = getattr(app, "hub_connect_url", "http://localhost:8888") - callback_url = ( - f"{hub_url.rstrip('/')}/api/beakerhub/internal/task-callback/" - f"{callback_token}" - ) try: - external_task_id = app.task_runner.submit_image_import( - node_image=node_image, - callback_url=callback_url, - callback_token=callback_token, + running_task = app.task_runner.submit( + ImageImportTask( + image=node_image.default_img_string, + entrypoint=("sh", "-c"), + command=("beaker context dump",), + node_image=node_image, + ) ) except Exception as error: task.status = "failed" @@ -66,13 +77,38 @@ def launch_import_task( 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 = external_task_id + 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", - external_task_id, + 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/services/test_service_wiring.py b/tests/unit/services/test_service_wiring.py index 5ccd471..204da91 100644 --- a/tests/unit/services/test_service_wiring.py +++ b/tests/unit/services/test_service_wiring.py @@ -7,7 +7,7 @@ from beakerhub.app import BeakerHub from beakerhub.services.dashboard.base import BaseDashboardService from beakerhub.services.task.base import BaseTaskRunnerService -from beakerhub.tasks.image_import.task import launch_import_task +from beakerhub.tasks.image_import.task import ImageImportTask, launch_import_task class DummyTaskRunnerService(BaseTaskRunnerService): @@ -42,11 +42,13 @@ def test_task_runner_receives_config_loaded_after_its_creation(): app = BeakerHub() runner = app.task_runner config = Config() - config.KubernetesTaskRunnerService.reporter_image = "registry.example/reporter:v1" + config.KubernetesTaskRunnerService.node_image_resources = { + "limits": {"cpu": "1"} + } app.update_config(config) - assert runner.reporter_image == "registry.example/reporter:v1" + assert runner.node_image_resources == {"limits": {"cpu": "1"}} def test_image_import_delegates_submission_to_task_runner(): @@ -54,11 +56,17 @@ def test_image_import_delegates_submission_to_task_runner(): 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_image_import.return_value = "external-task-id" + app.task_runner.submit.return_value = MagicMock(external_id="external-task-id") task = launch_import_task(db, app, node_image) - app.task_runner.submit_image_import.assert_called_once() + 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..fb42da0 --- /dev/null +++ b/tests/unit/services/test_task_base.py @@ -0,0 +1,57 @@ +"""Tests for task-runner lifecycle value objects.""" + +from dataclasses import dataclass + +import pytest + +from beakerhub.services.task.base import ( + BaseTaskRunnerService, + RunningTask, + TaskOutput, + TaskStatus, +) +from beakerhub.tasks.base import BaseTaskDefinition + + +@dataclass(frozen=True) +class ExampleTask(BaseTaskDefinition): + @property + def task_type(self) -> str: + return "example" + + +class CompleteTaskRunner(BaseTaskRunnerService): + def get_status(self, external_id: str) -> TaskStatus: + return TaskStatus("completed") + + def get_output(self, external_id: str) -> TaskOutput: + return TaskOutput("output", "") + + +def test_running_task_uses_its_external_id_for_runner_operations(): + runner = CompleteTaskRunner() + task = RunningTask("runtime-123", ExampleTask(), runner) + + assert task.status == TaskStatus("completed") + assert task.done is True + assert task.output == TaskOutput("output", "") + + +@pytest.mark.asyncio +async def test_running_task_returns_terminal_status(): + runner = CompleteTaskRunner() + task = RunningTask("runtime-123", ExampleTask(), runner) + + assert await task.await_completion() == TaskStatus("completed") + + +@pytest.mark.asyncio +async def test_running_task_raises_on_timeout(): + class PendingTaskRunner(CompleteTaskRunner): + def get_status(self, external_id: str) -> TaskStatus: + return TaskStatus("running") + + task = RunningTask("runtime-123", ExampleTask(), PendingTaskRunner()) + + with pytest.raises(TimeoutError, match="runtime-123"): + await task.await_completion(timeout=0) From e3ce23408dbc271d24a0e6d9ba7dfc4860bcbfb6 Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Fri, 28 Aug 2026 16:45:19 -0600 Subject: [PATCH 4/8] Finishing off ECS Task Runner --- .../services/task/aws_ecs_task_runner.py | 379 +++++++++++++++++- .../unit/services/test_aws_ecs_task_runner.py | 170 ++++++++ 2 files changed, 543 insertions(+), 6 deletions(-) create mode 100644 tests/unit/services/test_aws_ecs_task_runner.py diff --git a/src/beakerhub/services/task/aws_ecs_task_runner.py b/src/beakerhub/services/task/aws_ecs_task_runner.py index f774f6d..e0a7ced 100644 --- a/src/beakerhub/services/task/aws_ecs_task_runner.py +++ b/src/beakerhub/services/task/aws_ecs_task_runner.py @@ -1,21 +1,388 @@ """AWS ECS implementation skeleton for the task-runner service.""" +import typing +from typing import Any + +import boto3 +import traitlets +from botocore.client import BaseClient +from botocore.exceptions import ClientError + from beakerhub.services.task.base import ( BaseTaskRunnerService, RunningTask, + TaskOutput, TaskStatus, ) -from beakerhub.tasks.base import BaseTaskDefinition +from beakerhub.tasks.base import BaseImageTask, BaseTaskDefinition class AwsEcsTaskRunnerService(BaseTaskRunnerService): """Run BeakerHub background tasks as AWS ECS tasks.""" + boto_client: BaseClient = traitlets.Instance(klass=BaseClient, config=False) + logs_client: BaseClient = traitlets.Instance(klass=BaseClient, config=False) + + log_group: str = traitlets.Unicode( + default_value="", + help="CloudWatch Logs group that receives ECS task container output.", + config=True, + ) + log_stream_prefix: str = traitlets.Unicode( + default_value="beakerhub-task", + help="CloudWatch Logs stream prefix for ECS task containers.", + config=True, + ) + execution_role_arn: str = traitlets.Unicode( + default_value="", + help=( + "IAM role ARN used by ECS to pull images and write container logs. " + "This role needs ECR access for private images and " + "logs:CreateLogStream and logs:PutLogEvents for the log group." + ), + config=True, + ) + task_role_arn: str = traitlets.Unicode( + default_value="", + help=( + "IAM role ARN assumed by the task container. Grant it access to " + "AWS services required by the task." + ), + config=True, + ) + subnets: list[str] = traitlets.List( + traitlets.Unicode(), + default_value=[], + help="Subnets used by Fargate task awsvpc network configuration.", + config=True, + ) + security_groups: list[str] = traitlets.List( + traitlets.Unicode(), + default_value=[], + help="Security groups used by Fargate task awsvpc network configuration.", + config=True, + ) + assign_public_ip: bool = traitlets.Bool( + False, + help="Assign a public IP address to Fargate tasks.", + config=True, + ) + cpu_architecture: typing.Literal["X86_64", "ARM64"] = traitlets.Enum( + values=["X86_64", "ARM64"], + default_value="X86_64", + help="CPU architecture for task containers.", + config=True, + ) + task_definition_name: str = traitlets.Unicode( + default_value="beakerhub-tasks", + help="""Name of task definition "family". This is should not include the colon or revision number. Will be created if it does not exist.""", + config=True, + ) + cluster_name: str = traitlets.Unicode( + help="Name of ECS cluster to run task within", + config=True, + ) + task_group: str = traitlets.Unicode( + help="Optional. Name of task group to associate with the task.", + config=True, + ) + launch_type: typing.Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"] = traitlets.Enum( + values=["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], + default_value="FARGATE", + help="Method of launching ECS task", + config=True, + ) + launch_options: dict = traitlets.Dict( + help="Optional. If provided, overrides keyword arguments passed to boto3.client.run_task. See https://docs.aws.amazon.com/boto3/latest/reference/services/ecs/client/run_task.html", + config=True, + ) + task_overrides: dict = traitlets.Dict( + default_value=None, + allow_none=True, + help="Optional. ECS task TaskOverride values to be included. See https://docs.aws.amazon.com/AmazonECS/latest/APIReference/API_TaskOverride.html", + config=True, + ) + aws_region: str = traitlets.Unicode( + help="Optional. AWS Region cluster resides in.", + config=True, + ) + default_task_cpu: str = traitlets.Unicode( + default_value="4 vcpu", + help="Optional. Default CPU for task nodes. Used to create the task_definition if it does not exist. Does not override per-task runs.", + config=True, + ) + default_task_memory: str = traitlets.Unicode( + default_value="16GB", + help="Optional. Default memory for task nodes. Used to create the task_definition if it does not exist. Does not override per-task runs.", + config=True, + ) + default_task_ephemeral_storage: str = traitlets.Int( + default_value=50, + help="Optional. Default ephemeral storage for task nodes in GB. Used to create the task_definition if it does not exist. Does not override per-task runs.", + config=True, + ) + + @traitlets.default("boto_client") + def _default_boto_client(self): + return boto3.client("ecs", region_name=self.aws_region or None) + + @traitlets.default("logs_client") + def _default_logs_client(self): + return boto3.client("logs", region_name=self.aws_region or None) + + @traitlets.default('launch_options') + def _default_launch_options(self): + return {} + + @traitlets.default('aws_region') + def _default_aws_region(self): + session = boto3._get_default_session() + return (session.region_name if session else None) or "us-east-1" + def submit(self, task: BaseTaskDefinition) -> RunningTask: - raise NotImplementedError("AWS ECS task submission is not implemented") + """Register an image-specific definition and launch one ECS task.""" + if not isinstance(task, BaseImageTask): + raise ValueError( + "AwsEcsTaskRunnerService only supports image tasks, not " + f"{task.task_type!r}" + ) + + self._validate_configuration() + definition_arn = self._register_task_definition(task) + request: dict[str, Any] = { + "cluster": self.cluster_name, + "taskDefinition": definition_arn, + "overrides": self._build_overrides(task), + "tags": [ + {"key": "beakerhub-task", "value": "true"}, + {"key": "beakerhub-task-type", "value": task.task_type}, + ], + } + if self.task_group: + request["group"] = self.task_group + if self.launch_type == "FARGATE" and "networkConfiguration" not in self.launch_options: + request["networkConfiguration"] = self._network_configuration() + if "capacityProviderStrategy" not in self.launch_options: + request["launchType"] = self.launch_type + request.update(self.launch_options) + + try: + response = self.boto_client.run_task(**request) + except ClientError as error: + raise RuntimeError(f"ECS failed to submit 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 RunningTask(tasks[0]["taskArn"], task, self) + + def _register_task_definition(self, task: BaseImageTask) -> str: + """Register the image selected by a task as an ECS definition revision.""" + container: dict[str, Any] = { + "name": "task-container", + "image": task.image, + "essential": True, + } + if task.entrypoint: + container["entryPoint"] = list(task.entrypoint) + if task.command: + container["command"] = list(task.command) + if task.working_directory: + container["workingDirectory"] = task.working_directory + if task.environment: + container["environment"] = [ + {"name": name, "value": value} + for name, value in task.environment.items() + ] + if self.log_group: + container["logConfiguration"] = { + "logDriver": "awslogs", + "options": { + "awslogs-group": self.log_group, + "awslogs-region": self.aws_region, + "awslogs-stream-prefix": self.log_stream_prefix, + }, + } + + request: dict[str, Any] = { + "family": self.task_definition_name, + "containerDefinitions": [container], + "cpu": self.default_task_cpu, + "memory": self.default_task_memory, + "runtimePlatform": { + "cpuArchitecture": self.cpu_architecture, + "operatingSystemFamily": "LINUX", + }, + } + if self.launch_type == "FARGATE": + request["ephemeralStorage"] = { + "sizeInGiB": self.default_task_ephemeral_storage + } + if self.execution_role_arn: + request["executionRoleArn"] = self.execution_role_arn + if self.task_role_arn: + request["taskRoleArn"] = self.task_role_arn + if self.launch_type == "FARGATE": + request.update({ + "networkMode": "awsvpc", + "requiresCompatibilities": ["FARGATE"], + }) + + try: + response = self.boto_client.register_task_definition(**request) + except ClientError as error: + raise RuntimeError(f"ECS failed to register task definition: {error}") from error + return response["taskDefinition"]["taskDefinitionArn"] + + def _build_overrides(self, task: BaseImageTask) -> dict[str, Any]: + overrides = dict(self.task_overrides or {}) + containers: list[dict[str, Any]] = [] + task_container: dict[str, Any] = {"name": "task-container"} + environment: dict[str, str] = {} + for container in overrides.pop("containerOverrides", []): + container = dict(container) + if container.get("name") != "task-container": + containers.append(container) + continue + environment.update({ + item["name"]: item["value"] + for item in container.pop("environment", []) + }) + task_container.update(container) + environment.update(task.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 _network_configuration(self) -> dict[str, Any]: + return { + "awsvpcConfiguration": { + "subnets": list(self.subnets), + "securityGroups": list(self.security_groups), + "assignPublicIp": "ENABLED" if self.assign_public_ip else "DISABLED", + } + } + + def _validate_configuration(self) -> None: + if not self.cluster_name.strip(): + raise ValueError("AwsEcsTaskRunnerService.cluster_name must be configured") + if not self.log_group: + raise ValueError("AwsEcsTaskRunnerService.log_group must be configured") + if not self.aws_region: + raise ValueError("AwsEcsTaskRunnerService.aws_region must be configured") + if not self.execution_role_arn: + raise ValueError( + "AwsEcsTaskRunnerService.execution_role_arn 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( + "AwsEcsTaskRunnerService.subnets must be configured for Fargate tasks" + ) + elif self.default_task_ephemeral_storage: + raise ValueError( + "default_task_ephemeral_storage is supported only for Fargate tasks" + ) + if ( + self.launch_type == "MANAGED_INSTANCES" + and not self.launch_options.get("capacityProviderStrategy") + ): + raise ValueError( + "Managed Instances tasks require launch_options.capacityProviderStrategy" + ) + + def get_status(self, external_id: str) -> TaskStatus: + task = self._describe_task(external_id) + if task is None: + return TaskStatus("failed", f"ECS task {external_id} was not found") + status = task.get("lastStatus", "UNKNOWN") + if status in {"PROVISIONING", "PENDING", "ACTIVATING"}: + return TaskStatus("pending", f"ECS task is {status.lower()}") + if status in {"RUNNING", "DEACTIVATING", "DEPROVISIONING", "STOPPING"}: + return TaskStatus("running", f"ECS task is {status.lower()}") + if status != "STOPPED": + return TaskStatus("failed", f"ECS task has unexpected status {status!r}") + + container = self._task_container(task) + if container and container.get("exitCode") == 0: + return TaskStatus("completed", "ECS task completed successfully") + reason = (container or {}).get("reason") or task.get("stoppedReason") or "Unknown failure" + return TaskStatus("failed", f"ECS task failed ({reason})") + + def get_output(self, external_id: str) -> TaskOutput | None: + """Return task output from the configured CloudWatch Logs group.""" + if not self.log_group: + return None + task = self._describe_task(external_id) + if task is None: + return None + container = self._task_container(task) + task_id = external_id.split('/')[-1] + container_name = (container or {}).get('name', 'task-container') + stream_name = f"{self.log_stream_prefix}/{container_name}/{task_id}" + + messages: list[str] = [] + next_token: str | None = None + while True: + request: dict[str, Any] = { + "logGroupName": self.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 for ECS task {external_id}: {error}" + ) from error + messages.extend(event["message"] for event in response.get("events", [])) + token = response.get("nextForwardToken") + if not token or token == next_token: + break + next_token = token + return TaskOutput(stdout="\n".join(messages), stderr="") + + def delete(self, external_id: str) -> None: + try: + self.boto_client.stop_task( + cluster=self.cluster_name, + task=external_id, + reason="Completed BeakerHub task cleanup", + ) + except ClientError as error: + raise RuntimeError(f"ECS failed to stop task {external_id}: {error}") from error + + def reap_stale_tasks(self) -> None: + """Stale-task reconciliation needs persisted task ownership and is deferred.""" + self.log.warning("ECS stale-task reaping is not implemented") + # TODO: Clean stale task-definition revisions as part of reaping tasks - def get_status(self, task_id: str) -> TaskStatus: - raise NotImplementedError("AWS ECS task status polling is not implemented") + def _describe_task(self, external_id: str) -> dict[str, Any] | None: + try: + response = self.boto_client.describe_tasks( + cluster=self.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 delete(self, task_id: str) -> None: - raise NotImplementedError("AWS ECS task deletion is not implemented") + @staticmethod + def _task_container(task: dict[str, Any]) -> dict[str, Any] | None: + return next( + (item for item in task.get("containers", []) if item.get("name") == "task-container"), + None, + ) 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..4b4590c --- /dev/null +++ b/tests/unit/services/test_aws_ecs_task_runner.py @@ -0,0 +1,170 @@ +"""Tests for the ECS task-runner task-definition configuration.""" + +from dataclasses import dataclass +from unittest.mock import Mock + +import boto3 +import pytest + +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 test_task_definition_includes_configured_iam_roles(): + ecs = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + ecs.register_task_definition = Mock( + return_value={ + "taskDefinition": {"taskDefinitionArn": "task-definition-arn"} + } + ) + runner = AwsEcsTaskRunnerService( + boto_client=ecs, + log_group="/beakerhub/tasks", + execution_role_arn="arn:aws:iam::123456789012:role/ecs-task-execution", + task_role_arn="arn:aws:iam::123456789012:role/beakerhub-task", + ) + + assert runner._register_task_definition(ExampleImageTask(image="example:latest")) == ( + "task-definition-arn" + ) + + request = ecs.register_task_definition.call_args.kwargs + assert request["executionRoleArn"] == ( + "arn:aws:iam::123456789012:role/ecs-task-execution" + ) + assert request["taskRoleArn"] == "arn:aws:iam::123456789012:role/beakerhub-task" + + +def configured_runner(**overrides): + values = { + "cluster_name": "tasks", + "log_group": "/beakerhub/tasks", + "aws_region": "us-east-1", + "execution_role_arn": "arn:aws:iam::123456789012:role/ecs-task-execution", + "subnets": ["subnet-123"], + } + values.update(overrides) + return AwsEcsTaskRunnerService(**values) + + +def test_submit_uses_configured_fargate_networking(): + ecs = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + ecs.run_task = Mock(return_value={"tasks": [{"taskArn": "task-arn"}]}) + runner = configured_runner( + boto_client=ecs, + security_groups=["sg-123"], + assign_public_ip=True, + ) + runner._register_task_definition = Mock(return_value="task-definition-arn") + + runner.submit(ExampleImageTask(image="example:latest")) + + request = ecs.run_task.call_args.kwargs + assert request["networkConfiguration"] == { + "awsvpcConfiguration": { + "subnets": ["subnet-123"], + "securityGroups": ["sg-123"], + "assignPublicIp": "ENABLED", + } + } + + +def test_task_definition_uses_configured_cpu_architecture(): + ecs = boto3.client( + "ecs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + ecs.register_task_definition = Mock( + return_value={"taskDefinition": {"taskDefinitionArn": "task-definition-arn"}} + ) + runner = configured_runner(boto_client=ecs, cpu_architecture="ARM64") + + runner._register_task_definition(ExampleImageTask(image="example:latest")) + + request = ecs.register_task_definition.call_args.kwargs + assert request["runtimePlatform"]["cpuArchitecture"] == "ARM64" + + +def test_task_overrides_merge_the_task_container_environment(): + runner = configured_runner( + task_overrides={ + "containerOverrides": [ + { + "name": "task-container", + "command": ["configured-command"], + "environment": [ + {"name": "FROM_CONFIG", "value": "configured"}, + {"name": "SHARED", "value": "configured"}, + ], + }, + {"name": "sidecar", "command": ["sidecar-command"]}, + ] + } + ) + + overrides = runner._build_overrides( + ExampleImageTask( + image="example:latest", + environment={"SHARED": "task", "FROM_TASK": "task"}, + ) + ) + + assert overrides["containerOverrides"] == [ + {"name": "sidecar", "command": ["sidecar-command"]}, + { + "name": "task-container", + "command": ["configured-command"], + "environment": [ + {"name": "FROM_CONFIG", "value": "configured"}, + {"name": "SHARED", "value": "task"}, + {"name": "FROM_TASK", "value": "task"}, + ], + }, + ] + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"cluster_name": ""}, "cluster_name"), + ({"log_group": ""}, "log_group"), + ({"aws_region": ""}, "aws_region"), + ({"execution_role_arn": ""}, "execution_role_arn"), + ({"subnets": []}, "subnets"), + ( + {"launch_type": "EC2"}, + "default_task_ephemeral_storage", + ), + ( + { + "launch_type": "MANAGED_INSTANCES", + "default_task_ephemeral_storage": 0, + }, + "capacityProviderStrategy", + ), + ], +) +def test_configuration_validation_rejects_incompatible_settings(overrides, message): + runner = configured_runner(**overrides) + + with pytest.raises(ValueError, match=message): + runner._validate_configuration() From 37e74b277085089ccf983f82bd1ca39534529f4f Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Mon, 31 Aug 2026 11:24:14 -0600 Subject: [PATCH 5/8] Created base Runtime classes for running processes in containers --- src/beakerhub/runtimes/__init__.py | 21 + src/beakerhub/runtimes/aws_ecs.py | 578 +++++++++++++++++++++++++ src/beakerhub/runtimes/base.py | 143 ++++++ src/beakerhub/runtimes/kubernetes.py | 339 +++++++++++++++ tests/unit/runtimes/__init__.py | 0 tests/unit/runtimes/test_aws_ecs.py | 149 +++++++ tests/unit/runtimes/test_base.py | 62 +++ tests/unit/runtimes/test_kubernetes.py | 111 +++++ 8 files changed, 1403 insertions(+) create mode 100644 src/beakerhub/runtimes/__init__.py create mode 100644 src/beakerhub/runtimes/aws_ecs.py create mode 100644 src/beakerhub/runtimes/base.py create mode 100644 src/beakerhub/runtimes/kubernetes.py create mode 100644 tests/unit/runtimes/__init__.py create mode 100644 tests/unit/runtimes/test_aws_ecs.py create mode 100644 tests/unit/runtimes/test_base.py create mode 100644 tests/unit/runtimes/test_kubernetes.py 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..cd1251f --- /dev/null +++ b/src/beakerhub/runtimes/aws_ecs.py @@ -0,0 +1,578 @@ +"""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. +""" + +from dataclasses import dataclass, field +from typing import Any, Literal, Mapping + +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 + + +@dataclass(frozen=True, kw_only=True) +class AwsEcsDefinition(BaseDefinition): + """An ECS workload definition. + + Set ``task_definition`` 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. + """ + + task_definition: str | None = None + container_name: str = "task-container" + cpu: str | None = None + memory: str | None = None + ephemeral_storage: int | None = None + tags: Mapping[str, str] = field(default_factory=dict) + + +class AwsEcsProcess(BaseProcess): + """A single launched ECS workload.""" + + runtime = traitlets.Instance(AwsEcsRuntime, allow_none=False) + + # Process traits remain configurable as service-specific overrides. Their + # defaults are resolved from the runtime after its configuration is loaded. + log_group: str = traitlets.Unicode(config=True) + log_stream_prefix: str = traitlets.Unicode(config=True) + execution_role_arn: str = traitlets.Unicode(config=True) + task_role_arn: str = traitlets.Unicode(config=True) + subnets: list[str] = traitlets.List(traitlets.Unicode(), config=True) + security_groups: list[str] = traitlets.List(traitlets.Unicode(), config=True) + assign_public_ip: bool = traitlets.Bool(config=True) + cpu_architecture: Literal["X86_64", "ARM64"] = traitlets.Enum( + values=["X86_64", "ARM64"], + config=True, + ) + task_definition_name: str = traitlets.Unicode(config=True) + cluster_name: str = traitlets.Unicode(config=True) + task_group: str = traitlets.Unicode(config=True) + launch_type: Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"] = ( + traitlets.Enum( + values=["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], + config=True, + ) + ) + launch_options: dict[str, Any] = traitlets.Dict(config=True) + task_overrides: dict[str, Any] | None = traitlets.Dict( + allow_none=True, + config=True, + ) + @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 + ) + + @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) + self._validate_configuration() + task_definition = definition.task_definition or self._register_task_definition() + external_id = self.runtime.run_task(self._run_task_request(task_definition)) + self.external_id = external_id + return self + + def _run_task_request(self, task_definition: str) -> dict[str, Any]: + definition = self._definition + request: dict[str, Any] = { + "cluster": self.cluster_name, + "taskDefinition": task_definition, + "overrides": self._build_overrides(), + "tags": self._tags(), + } + if self.task_group: + request["group"] = self.task_group + if self.launch_type == "FARGATE" and "networkConfiguration" not in self.launch_options: + request["networkConfiguration"] = self.runtime.awsvpc_network_configuration( + self.subnets, + self.security_groups, + self.assign_public_ip, + ) + if "capacityProviderStrategy" not in self.launch_options: + request["launchType"] = self.launch_type + request.update(self.launch_options) + return request + + @property + def _definition(self) -> AwsEcsDefinition: + if not isinstance(self.definition, AwsEcsDefinition): + raise TypeError("AwsEcsProcess requires an AwsEcsDefinition") + return self.definition + + def _tags(self) -> list[dict[str, str]]: + tags = dict(self._definition.tags) + tags.setdefault("beakerhub-process", "true") + tags.setdefault("beakerhub-process-type", self.process_type) + return [{"key": key, "value": value} for key, value in tags.items()] + + def _register_task_definition(self) -> str: + definition = self._definition + if not definition.image: + raise ValueError("AwsEcsDefinition.image is required without task_definition") + container: dict[str, Any] = { + "name": definition.container_name, + "image": definition.image, + "essential": True, + } + if definition.entrypoint: + container["entryPoint"] = list(definition.entrypoint) + if definition.command: + container["command"] = list(definition.command) + if definition.working_directory: + container["workingDirectory"] = definition.working_directory + if definition.environment: + container["environment"] = [ + {"name": name, "value": value} + for name, value in definition.environment.items() + ] + 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, + }, + } + + request: dict[str, Any] = { + "family": self.task_definition_name, + "containerDefinitions": [container], + "cpu": definition.cpu or self.runtime.default_task_cpu, + "memory": definition.memory or self.runtime.default_task_memory, + "runtimePlatform": { + "cpuArchitecture": self.cpu_architecture, + "operatingSystemFamily": "LINUX", + }, + } + if self.launch_type == "FARGATE": + request.update( + { + "ephemeralStorage": { + "sizeInGiB": self._ephemeral_storage + }, + "networkMode": "awsvpc", + "requiresCompatibilities": ["FARGATE"], + } + ) + if self.execution_role_arn: + request["executionRoleArn"] = self.execution_role_arn + if self.task_role_arn: + request["taskRoleArn"] = self.task_role_arn + + try: + response = self.runtime.ecs_client.register_task_definition(**request) + except ClientError as error: + raise RuntimeError(f"ECS failed to register task definition: {error}") from error + return response["taskDefinition"]["taskDefinitionArn"] + + def _build_overrides(self) -> dict[str, Any]: + overrides = dict(self.task_overrides or {}) + containers: list[dict[str, Any]] = [] + task_container: dict[str, Any] = {"name": self._definition.container_name} + environment: dict[str, str] = {} + for container in overrides.pop("containerOverrides", []): + container = dict(container) + if container.get("name") != self._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(self._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 _validate_configuration(self) -> None: + if not self.cluster_name.strip(): + raise ValueError("AwsEcsProcess.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("AwsEcsProcess.subnets must be configured for Fargate workloads") + elif self._ephemeral_storage and not self._definition.task_definition: + 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" + ) + + @property + def _ephemeral_storage(self) -> int: + value = self._definition.ephemeral_storage + return value if value is not None else self.runtime.default_task_ephemeral_storage + + 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.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) + if container and container.get("exitCode") == 0: + return ProcessStatus("completed", "ECS task completed successfully") + reason = ( + (container or {}).get("reason") + or task.get("stoppedReason") + or "Unknown failure" + ) + return ProcessStatus("failed", f"ECS task failed ({reason})") + + def collect_output(self) -> ProcessOutput | None: + """Return retained CloudWatch output for this ECS workload.""" + if not self.log_group or not self.external_id: + return None + task = self.runtime.describe_task(self.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 = f"{self.log_stream_prefix}/{container_name}/{task_id}" + return ProcessOutput( + stdout="\n".join(self.runtime.get_log_events(self.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.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, + ) + + +# 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..3888ff1 --- /dev/null +++ b/src/beakerhub/runtimes/base.py @@ -0,0 +1,143 @@ +"""Provider-neutral runtime workload contracts and composition helpers.""" + +import uuid +from dataclasses import dataclass, field +from typing import Any, Literal, Mapping + +from traitlets import Instance, Type, default +from traitlets.config import LoggingConfigurable + + +ProcessState = Literal["pending", "running", "completed", "failed"] +ProcessType = Literal["task", "service"] + + +@dataclass(frozen=True) +class ProcessStatus: + """The runtime's current view of a launched process.""" + + state: ProcessState + message: str | 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 + + +@dataclass(frozen=True, kw_only=True) +class BaseDefinition: + """Provider-neutral description of an OCI workload. + + A definition describes the workload to launch, not its runtime identity or + lifecycle. Provider definitions may add scheduling and launch fields. + """ + + image: str | None = None + entrypoint: tuple[str, ...] = () + command: tuple[str, ...] = () + working_directory: str | None = None + environment: Mapping[str, str] = field(default_factory=dict) + labels: Mapping[str, str] = field(default_factory=dict) + + +class BaseRuntime(LoggingConfigurable): + """Own a provider connection and provider-wide runtime defaults.""" + + +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: + self.id = kwargs.pop("id", uuid.uuid4().hex) + self.process_type = process_type + self.external_id = external_id + self.definition = definition + super().__init__(**kwargs) + + @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 + + +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, + ) + runtime = Instance(BaseRuntime, allow_none=False) + + @default("runtime") + def _default_runtime(self) -> BaseRuntime: + return self.runtime_class(parent=self) + + def create_definition(self, **kwargs: Any) -> BaseDefinition: + """Create a definition using this bundle's provider implementation.""" + return self.definition_class(**kwargs) + + def start_process( + self, + definition: BaseDefinition, + *, + process_type: ProcessType = "task", + **kwargs: Any, + ) -> BaseProcess: + """Launch a provider process using this bundle's shared runtime.""" + if not isinstance(definition, self.definition_class): + raise TypeError( + f"{self.__class__.__name__} requires " + f"{self.definition_class.__name__}, not {type(definition).__name__}" + ) + return self.process_class.start( + definition, + process_type=process_type, + runtime=self.runtime, + parent=self, + **kwargs, + ) diff --git a/src/beakerhub/runtimes/kubernetes.py b/src/beakerhub/runtimes/kubernetes.py new file mode 100644 index 0000000..577e54d --- /dev/null +++ b/src/beakerhub/runtimes/kubernetes.py @@ -0,0 +1,339 @@ +"""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 dataclasses import dataclass, field +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.", + ) + + @staticmethod + def get_clients() -> tuple[k8s_client.BatchV1Api, k8s_client.CoreV1Api]: + """Load Kubernetes configuration and create the required API clients.""" + try: + k8s_config.load_incluster_config() + except k8s_config.ConfigException: + k8s_config.load_kube_config() + return k8s_client.BatchV1Api(), k8s_client.CoreV1Api() + + @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 + + +@dataclass(frozen=True, kw_only=True) +class KubernetesDefinition(BaseDefinition): + """A Kubernetes Job workload definition.""" + + namespace: str | None = None + resources: Mapping[str, Any] = field(default_factory=dict) + node_selector: Mapping[str, str] | None = None + tolerations: tuple[Mapping[str, Any], ...] | None = None + service_account: str | None = None + backoff_limit: int = 0 + active_deadline_seconds: int | None = 300 + ttl_seconds_after_finished: int | None = 600 + name_prefix: str = "beaker-process" + + +class KubernetesProcess(BaseProcess): + """A Kubernetes Job-backed process.""" + + runtime = Instance(KubernetesRuntime, allow_none=False) + + def __init__(self, definition: KubernetesDefinition, **kwargs: Any) -> None: + super().__init__(definition, **kwargs) + + @property + def _definition(self) -> KubernetesDefinition: + if not isinstance(self.definition, KubernetesDefinition): + raise TypeError("KubernetesProcess requires a KubernetesDefinition") + return self.definition + + @property + def namespace(self) -> str: + return self._definition.namespace or self.runtime.namespace + + @property + def node_selector(self) -> Mapping[str, str]: + if self._definition.node_selector is not None: + return self._definition.node_selector + return self.runtime.node_selector + + @property + def tolerations(self) -> tuple[Mapping[str, Any], ...] | list[dict[str, Any]]: + if self._definition.tolerations is not None: + return self._definition.tolerations + return self.runtime.tolerations + + @property + def service_account(self) -> str: + if self._definition.service_account is not None: + return self._definition.service_account + return self.runtime.service_account + + @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() + batch_api, _ = self.runtime.get_clients() + batch_api.create_namespaced_job(namespace=self.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(self.node_selector) or None, + tolerations=( + [k8s_client.V1Toleration(**item) for item in self.tolerations] + if self.tolerations + else None + ), + service_account_name=self.service_account or None, + ) + return k8s_client.V1Job( + api_version="batch/v1", + kind="Job", + metadata=k8s_client.V1ObjectMeta( + name=self.external_id, + namespace=self.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") + batch_api, core_api = self.runtime.get_clients() + try: + job = batch_api.read_namespaced_job( + name=self.external_id, + namespace=self.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(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.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 + _, core_api = self.runtime.get_clients() + pods = core_api.list_namespaced_pod( + namespace=self.namespace, + label_selector=f"job-name={self.external_id}", + ) + if not pods.items: + return None + response = core_api.read_namespaced_pod_log( + name=pods.items[0].metadata.name, + namespace=self.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 + batch_api, _ = self.runtime.get_clients() + try: + batch_api.delete_namespaced_job( + name=self.external_id, + namespace=self.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, + ) + + +# 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/tests/unit/runtimes/__init__.py b/tests/unit/runtimes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/runtimes/test_aws_ecs.py b/tests/unit/runtimes/test_aws_ecs.py new file mode 100644 index 0000000..c48900d --- /dev/null +++ b/tests/unit/runtimes/test_aws_ecs.py @@ -0,0 +1,149 @@ +"""Tests for ECS runtime process construction and control-plane behavior.""" + +from unittest.mock import Mock + +import boto3 + +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_process_uses_runtime_defaults_and_allows_process_overrides(): + runtime = ecs_runtime() + definition = AwsEcsDefinition(task_definition="notebook:1") + + process = AwsEcsProcess(definition, runtime=runtime) + overridden = AwsEcsProcess( + definition, + runtime=runtime, + cluster_name="process-cluster", + subnets=["subnet-process"], + ) + + assert process.cluster_name == "runtime-cluster" + assert process.subnets == ["subnet-runtime"] + assert process.security_groups == ["sg-runtime"] + assert process.log_group == "/beakerhub/runtime" + assert process.launch_options == {"enableExecuteCommand": True} + assert overridden.cluster_name == "process-cluster" + assert overridden.subnets == ["subnet-process"] + + process.subnets.append("local-change") + assert runtime.subnets == ["subnet-runtime"] + + +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( + task_definition="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_resources_override_runtime_defaults(): + runtime = ecs_runtime( + default_task_cpu="2 vcpu", + default_task_memory="8GB", + default_task_ephemeral_storage=40, + ) + runtime.ecs_client.register_task_definition = Mock( + return_value={"taskDefinition": {"taskDefinitionArn": "definition-arn"}} + ) + process = AwsEcsProcess( + AwsEcsDefinition( + image="example:latest", + cpu="1 vcpu", + memory="2GB", + ephemeral_storage=25, + ), + runtime=runtime, + ) + + assert process._register_task_definition() == "definition-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_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..3086eaf --- /dev/null +++ b/tests/unit/runtimes/test_base.py @@ -0,0 +1,62 @@ +"""Tests for provider-neutral runtime contracts.""" + +from dataclasses import dataclass + +from beakerhub.runtimes.base import ( + BaseDefinition, + BaseProcess, + BaseRuntime, + BaseRuntimeBundle, + ProcessStatus, +) + + +@dataclass(frozen=True, kw_only=True) +class ExampleDefinition(BaseDefinition): + value: str = "example" + + +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 + + +def test_bundle_creates_definition_and_process_with_shared_runtime(): + bundle = ExampleBundle() + + definition = bundle.create_definition(value="configured") + process = bundle.start_process(definition) + + assert isinstance(definition, ExampleDefinition) + assert definition.value == "configured" + assert isinstance(process.runtime, ExampleRuntime) + assert process.runtime is bundle.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 diff --git a/tests/unit/runtimes/test_kubernetes.py b/tests/unit/runtimes/test_kubernetes.py new file mode 100644 index 0000000..b746748 --- /dev/null +++ b/tests/unit/runtimes/test_kubernetes.py @@ -0,0 +1,111 @@ +"""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() + monkeypatch.setattr(instance, "get_clients", lambda: (batch_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() + monkeypatch.setattr(instance, "get_clients", lambda: (batch_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" From e40c2abef91d8f99ed435be6d57ddf1d5f51a4b7 Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Tue, 8 Sep 2026 10:49:27 -0600 Subject: [PATCH 6/8] Bulk commit finishing off ECS spawner work --- .gitignore | 1 + src/beakerhub/app.py | 68 +- src/beakerhub/runtimes/aws_ecs.py | 465 ++++++++----- src/beakerhub/runtimes/base.py | 136 +++- src/beakerhub/runtimes/kubernetes.py | 140 ++-- .../services/dashboard/aws_ecs_dashboard.py | 312 ++++++++- src/beakerhub/services/dashboard/base.py | 7 +- src/beakerhub/services/dashboard/handlers.py | 52 +- .../dashboard/kubernetes_dashboard.py | 119 +++- .../services/spawner/aws_ecs_spawner.py | 642 +++++++++++++++--- src/beakerhub/services/spawner/base.py | 89 ++- .../services/spawner/kubernetes_spawner.py | 87 +-- .../services/task/aws_ecs_task_runner.py | 516 ++++---------- src/beakerhub/services/task/base.py | 105 +-- src/beakerhub/services/task/handlers.py | 10 +- .../services/task/kubernetes_task_runner.py | 256 ++----- src/beakerhub/tasks/image_import/task.py | 49 +- tests/unit/runtimes/test_aws_ecs.py | 176 ++++- tests/unit/runtimes/test_base.py | 25 +- tests/unit/runtimes/test_kubernetes.py | 6 +- tests/unit/services/spawner/test_aws_ecs.py | 587 +++++++++++++--- .../unit/services/spawner/test_kubernetes.py | 97 ++- tests/unit/services/test_aws_ecs_dashboard.py | 175 +++++ .../unit/services/test_aws_ecs_task_runner.py | 205 +++--- .../unit/services/test_dashboard_handlers.py | 36 + .../services/test_kubernetes_task_runner.py | 62 ++ tests/unit/services/test_service_wiring.py | 75 ++ tests/unit/services/test_task_base.py | 93 ++- tests/unit/tasks/test_image_import.py | 55 ++ ...{PodLogViewer.vue => SessionLogViewer.vue} | 20 +- ui/src/pages/Dashboard.vue | 2 +- ui/src/pages/admin/AdminDashboard.vue | 548 +++++---------- ui/src/pages/admin/AdminSessions.vue | 10 +- ui/src/stores/admin.ts | 88 +-- 34 files changed, 3425 insertions(+), 1889 deletions(-) create mode 100644 tests/unit/services/test_aws_ecs_dashboard.py create mode 100644 tests/unit/services/test_dashboard_handlers.py create mode 100644 tests/unit/services/test_kubernetes_task_runner.py create mode 100644 tests/unit/tasks/test_image_import.py rename ui/src/components/admin/{PodLogViewer.vue => SessionLogViewer.vue} (92%) 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/src/beakerhub/app.py b/src/beakerhub/app.py index 5c8f3ad..1b1f0f3 100644 --- a/src/beakerhub/app.py +++ b/src/beakerhub/app.py @@ -13,6 +13,8 @@ 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.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 @@ -37,8 +39,28 @@ class BeakerHub(JupyterHub): description = "Beakerhub version of: \n" + str(JupyterHub.description) example = "Beakerhub version of: \n" + str(JupyterHub.examples) + runtime_bundle_class = Type( + allow_none=True, + klass=BaseRuntimeBundle, + config=True, + help="Configures a default suite of runtime based configuration options." + ) + runtime_bundle = Instance( + klass=BaseRuntimeBundle, + allow_none=True, + ) + + runtime_class = Type( + klass=BaseRuntime, + config=True, + help="Configures a default suite of runtime based configuration options." + ) + runtime = Instance( + klass=BaseRuntime, + allow_none=True, + ) + task_runner_class = Type( - KubernetesTaskRunnerService, klass=BaseTaskRunnerService, config=True, help="Task-runner service used for background workloads.", @@ -47,8 +69,8 @@ class BeakerHub(JupyterHub): BaseTaskRunnerService, allow_none=False, ) + dashboard_service_class = Type( - KubernetesDashboardService, klass=BaseDashboardService, config=True, help="Service used to collect runtime dashboard data and session logs.", @@ -110,15 +132,49 @@ 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.services.spawner.kubernetes_spawner 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) @@ -132,6 +188,10 @@ def update_config(self, config): 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: diff --git a/src/beakerhub/runtimes/aws_ecs.py b/src/beakerhub/runtimes/aws_ecs.py index cd1251f..f456e92 100644 --- a/src/beakerhub/runtimes/aws_ecs.py +++ b/src/beakerhub/runtimes/aws_ecs.py @@ -5,8 +5,9 @@ policies while delegating provider operations to these classes. """ -from dataclasses import dataclass, field -from typing import Any, Literal, Mapping +import hashlib +import re +from typing import Any, Literal import boto3 import traitlets @@ -218,56 +219,90 @@ def get_log_events( 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 + -@dataclass(frozen=True, kw_only=True) class AwsEcsDefinition(BaseDefinition): """An ECS workload definition. - Set ``task_definition`` for a deployment-managed ECS definition, as session + 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. """ - task_definition: str | None = None - container_name: str = "task-container" - cpu: str | None = None - memory: str | None = None - ephemeral_storage: int | None = None - tags: Mapping[str, str] = field(default_factory=dict) + 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 -class AwsEcsProcess(BaseProcess): - """A single launched ECS workload.""" + @traitlets.default("memory") + def _default_memory(self) -> str: + return self.runtime.default_task_memory - runtime = traitlets.Instance(AwsEcsRuntime, allow_none=False) + @traitlets.default("ephemeral_storage") + def _default_ephemeral_storage(self) -> int: + return self.runtime.default_task_ephemeral_storage - # Process traits remain configurable as service-specific overrides. Their - # defaults are resolved from the runtime after its configuration is loaded. - log_group: str = traitlets.Unicode(config=True) - log_stream_prefix: str = traitlets.Unicode(config=True) - execution_role_arn: str = traitlets.Unicode(config=True) - task_role_arn: str = traitlets.Unicode(config=True) - subnets: list[str] = traitlets.List(traitlets.Unicode(), config=True) - security_groups: list[str] = traitlets.List(traitlets.Unicode(), config=True) - assign_public_ip: bool = traitlets.Bool(config=True) - cpu_architecture: Literal["X86_64", "ARM64"] = traitlets.Enum( - values=["X86_64", "ARM64"], - config=True, - ) - task_definition_name: str = traitlets.Unicode(config=True) - cluster_name: str = traitlets.Unicode(config=True) - task_group: str = traitlets.Unicode(config=True) - launch_type: Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"] = ( - traitlets.Enum( - values=["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], - config=True, - ) - ) - launch_options: dict[str, Any] = traitlets.Dict(config=True) - task_overrides: dict[str, Any] | None = traitlets.Dict( - allow_none=True, - config=True, - ) @traitlets.default("log_group") def _default_log_group(self) -> str: return self.runtime.log_group @@ -302,7 +337,17 @@ def _default_cpu_architecture(self) -> Literal["X86_64", "ARM64"]: @traitlets.default("task_definition_name") def _default_task_definition_name(self) -> str: - return self.runtime.task_definition_name + 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: @@ -330,74 +375,124 @@ def _default_task_overrides(self) -> dict[str, Any] | None: else None ) - @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) - self._validate_configuration() - task_definition = definition.task_definition or self._register_task_definition() - external_id = self.runtime.run_task(self._run_task_request(task_definition)) - self.external_id = external_id - return self + 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 _run_task_request(self, task_definition: str) -> dict[str, Any]: - definition = self._definition - request: dict[str, Any] = { - "cluster": self.cluster_name, - "taskDefinition": task_definition, - "overrides": self._build_overrides(), - "tags": self._tags(), - } - if self.task_group: - request["group"] = self.task_group - if self.launch_type == "FARGATE" and "networkConfiguration" not in self.launch_options: - request["networkConfiguration"] = self.runtime.awsvpc_network_configuration( - self.subnets, - self.security_groups, - self.assign_public_ip, + 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 ) - if "capacityProviderStrategy" not in self.launch_options: - request["launchType"] = self.launch_type - request.update(self.launch_options) - return request + 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 - @property - def _definition(self) -> AwsEcsDefinition: - if not isinstance(self.definition, AwsEcsDefinition): - raise TypeError("AwsEcsProcess requires an AwsEcsDefinition") - return self.definition + @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 _tags(self) -> list[dict[str, str]]: - tags = dict(self._definition.tags) - tags.setdefault("beakerhub-process", "true") - tags.setdefault("beakerhub-process-type", self.process_type) - return [{"key": key, "value": value} for key, value in tags.items()] + 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) - def _register_task_definition(self) -> str: - definition = self._definition - if not definition.image: - raise ValueError("AwsEcsDefinition.image is required without task_definition") + @property + def aws_task_definition(self): container: dict[str, Any] = { - "name": definition.container_name, - "image": definition.image, + "name": self.container_name, + "image": self.image, "essential": True, } - if definition.entrypoint: - container["entryPoint"] = list(definition.entrypoint) - if definition.command: - container["command"] = list(definition.command) - if definition.working_directory: - container["workingDirectory"] = definition.working_directory - if definition.environment: + 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 definition.environment.items() + 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"] = { @@ -409,52 +504,111 @@ def _register_task_definition(self) -> str: }, } - request: dict[str, Any] = { + definition_dict: dict[str, Any] = { "family": self.task_definition_name, - "containerDefinitions": [container], - "cpu": definition.cpu or self.runtime.default_task_cpu, - "memory": definition.memory or self.runtime.default_task_memory, + "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": - request.update( + definition_dict.update( { "ephemeralStorage": { - "sizeInGiB": self._ephemeral_storage + "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: - request["executionRoleArn"] = self.execution_role_arn + definition_dict["executionRoleArn"] = self.execution_role_arn if self.task_role_arn: - request["taskRoleArn"] = self.task_role_arn + definition_dict["taskRoleArn"] = self.task_role_arn + return definition_dict - try: - response = self.runtime.ecs_client.register_task_definition(**request) - except ClientError as error: - raise RuntimeError(f"ECS failed to register task definition: {error}") from error - return response["taskDefinition"]["taskDefinitionArn"] +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]: - overrides = dict(self.task_overrides or {}) + definition = self.definition + overrides = dict(definition.task_overrides or {}) containers: list[dict[str, Any]] = [] - task_container: dict[str, Any] = {"name": self._definition.container_name} + 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") != self._definition.container_name: + 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(self._definition.environment) + environment.update(definition.environment) if environment: task_container["environment"] = [ {"name": name, "value": value} for name, value in environment.items() @@ -463,40 +617,11 @@ def _build_overrides(self) -> dict[str, Any]: overrides["containerOverrides"] = containers return overrides - def _validate_configuration(self) -> None: - if not self.cluster_name.strip(): - raise ValueError("AwsEcsProcess.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("AwsEcsProcess.subnets must be configured for Fargate workloads") - elif self._ephemeral_storage and not self._definition.task_definition: - 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" - ) - - @property - def _ephemeral_storage(self) -> int: - value = self._definition.ephemeral_storage - return value if value is not None else self.runtime.default_task_ephemeral_storage - 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.cluster_name, self.external_id) + 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") @@ -508,28 +633,40 @@ def describe(self) -> ProcessStatus: return ProcessStatus("failed", f"ECS task has unexpected status {status!r}") container = self._task_container(task) - if container and container.get("exitCode") == 0: - return ProcessStatus("completed", "ECS task completed successfully") + 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})") + 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.log_group or not self.external_id: + if not self.definition.log_group or not self.external_id: return None - task = self.runtime.describe_task(self.cluster_name, self.external_id) + 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 = f"{self.log_stream_prefix}/{container_name}/{task_id}" + 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.log_group, stream_name)), + stdout="\n".join(self.runtime.get_log_events(self.definition.log_group, stream_name)), stderr="", ) @@ -538,7 +675,7 @@ def stop(self) -> None: if not self.external_id: return self.runtime.stop_task( - self.cluster_name, + self.definition.cluster_name, self.external_id, f"Stopped BeakerHub {self.process_type} cleanup", ) @@ -548,7 +685,7 @@ def _task_container(self, task: dict[str, Any]) -> dict[str, Any] | None: ( item for item in task.get("containers", []) - if item.get("name") == self._definition.container_name + if item.get("name") == self.definition.container_name ), None, ) @@ -573,6 +710,22 @@ class AwsEcsRuntimeBundle(BaseRuntimeBundle): 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 index 3888ff1..1365620 100644 --- a/src/beakerhub/runtimes/base.py +++ b/src/beakerhub/runtimes/base.py @@ -1,15 +1,17 @@ """Provider-neutral runtime workload contracts and composition helpers.""" +import asyncio import uuid -from dataclasses import dataclass, field -from typing import Any, Literal, Mapping +from dataclasses import dataclass +from time import monotonic +from typing import Any, Literal, TypeAlias -from traitlets import Instance, Type, default +from traitlets import Dict, Instance, List, Type, Unicode, default from traitlets.config import LoggingConfigurable -ProcessState = Literal["pending", "running", "completed", "failed"] -ProcessType = Literal["task", "service"] +ProcessState: TypeAlias = Literal["pending", "running", "completed", "failed"] +ProcessType: TypeAlias = Literal["task", "service"] @dataclass(frozen=True) @@ -18,6 +20,7 @@ class ProcessStatus: state: ProcessState message: str | None = None + exit_code: int | None = None @property def done(self) -> bool: @@ -33,24 +36,37 @@ class ProcessOutput: stderr: str -@dataclass(frozen=True, kw_only=True) -class BaseDefinition: +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. Provider definitions may add scheduling and launch fields. + lifecycle. Its unset provider fields resolve from ``runtime``. Provider + definitions may add scheduling and launch fields. """ - image: str | None = None - entrypoint: tuple[str, ...] = () - command: tuple[str, ...] = () - working_directory: str | None = None - environment: Mapping[str, str] = field(default_factory=dict) - labels: Mapping[str, str] = field(default_factory=dict) - - -class BaseRuntime(LoggingConfigurable): - """Own a provider connection and provider-wide runtime defaults.""" + 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): @@ -66,11 +82,28 @@ def __init__( 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( @@ -100,6 +133,20 @@ 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.""" @@ -111,33 +158,70 @@ class BaseRuntimeBundle(LoggingConfigurable): default_value=BaseDefinition, config=True, ) - runtime = Instance(BaseRuntime, allow_none=False) + # 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 + ) - @default("runtime") - def _default_runtime(self) -> BaseRuntime: - return self.runtime_class(parent=self) + @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, **kwargs: Any) -> BaseDefinition: + 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 this bundle's shared runtime.""" + """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=self.runtime, - parent=self, + runtime=runtime, + parent=runtime, **kwargs, ) diff --git a/src/beakerhub/runtimes/kubernetes.py b/src/beakerhub/runtimes/kubernetes.py index 577e54d..6072acd 100644 --- a/src/beakerhub/runtimes/kubernetes.py +++ b/src/beakerhub/runtimes/kubernetes.py @@ -5,7 +5,6 @@ proxy readiness until it is explicitly adapted to this runtime layer. """ -from dataclasses import dataclass, field from typing import Any, Mapping from uuid import uuid4 @@ -59,14 +58,18 @@ class KubernetesRuntime(BaseRuntime): help="Labels applied to every runtime workload.", ) - @staticmethod - def get_clients() -> tuple[k8s_client.BatchV1Api, k8s_client.CoreV1Api]: - """Load Kubernetes configuration and create the required API clients.""" + 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() - return k8s_client.BatchV1Api(), k8s_client.CoreV1Api() + self.batch_api = k8s_client.BatchV1Api() + self.core_api = k8s_client.CoreV1Api() + super().__init__(**kwargs) @staticmethod def build_resource_requirements( @@ -101,56 +104,47 @@ def container_failure_message(status: Any, label: str) -> str | None: return None -@dataclass(frozen=True, kw_only=True) class KubernetesDefinition(BaseDefinition): """A Kubernetes Job workload definition.""" - namespace: str | None = None - resources: Mapping[str, Any] = field(default_factory=dict) - node_selector: Mapping[str, str] | None = None - tolerations: tuple[Mapping[str, Any], ...] | None = None - service_account: str | None = None - backoff_limit: int = 0 - active_deadline_seconds: int | None = 300 - ttl_seconds_after_finished: int | None = 600 - name_prefix: str = "beaker-process" + 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: - super().__init__(definition, **kwargs) - - @property - def _definition(self) -> KubernetesDefinition: - if not isinstance(self.definition, KubernetesDefinition): + if not isinstance(definition, KubernetesDefinition): raise TypeError("KubernetesProcess requires a KubernetesDefinition") - return self.definition - - @property - def namespace(self) -> str: - return self._definition.namespace or self.runtime.namespace - - @property - def node_selector(self) -> Mapping[str, str]: - if self._definition.node_selector is not None: - return self._definition.node_selector - return self.runtime.node_selector - - @property - def tolerations(self) -> tuple[Mapping[str, Any], ...] | list[dict[str, Any]]: - if self._definition.tolerations is not None: - return self._definition.tolerations - return self.runtime.tolerations - - @property - def service_account(self) -> str: - if self._definition.service_account is not None: - return self._definition.service_account - return self.runtime.service_account + super().__init__(definition, **kwargs) @classmethod def start( @@ -167,17 +161,19 @@ def start( ) self = cls(definition, process_type=process_type, **kwargs) self.external_id = self._job_name() - batch_api, _ = self.runtime.get_clients() - batch_api.create_namespaced_job(namespace=self.namespace, body=self._job()) + 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" + prefix = self.definition.name_prefix.rstrip("-") or "beaker-process" return f"{prefix}-{uuid4().hex[:8]}" def _job(self) -> k8s_client.V1Job: - definition = self._definition + definition = self.definition labels = { **self.runtime.base_labels, **definition.labels, @@ -201,20 +197,20 @@ def _job(self) -> k8s_client.V1Job: pod_spec = k8s_client.V1PodSpec( containers=[container], restart_policy="Never", - node_selector=dict(self.node_selector) or None, + node_selector=dict(definition.node_selector) or None, tolerations=( - [k8s_client.V1Toleration(**item) for item in self.tolerations] - if self.tolerations + [k8s_client.V1Toleration(**item) for item in definition.tolerations] + if definition.tolerations else None ), - service_account_name=self.service_account or 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=self.namespace, + namespace=definition.namespace, labels=labels, ), spec=k8s_client.V1JobSpec( @@ -232,11 +228,10 @@ 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") - batch_api, core_api = self.runtime.get_clients() try: - job = batch_api.read_namespaced_job( + job = self.runtime.batch_api.read_namespaced_job( name=self.external_id, - namespace=self.namespace, + namespace=self.definition.namespace, ) except k8s_client.ApiException as error: if error.status == 404: @@ -249,7 +244,7 @@ def describe(self) -> ProcessStatus: if status.failed and status.failed > 0: return ProcessStatus( "failed", - self._failure_message(core_api), + self._failure_message(self.runtime.core_api), ) if status.active and status.active > 0: return ProcessStatus("running", "Job is running") @@ -258,7 +253,7 @@ def describe(self) -> ProcessStatus: def _failure_message(self, core_api: k8s_client.CoreV1Api) -> str: try: pods = core_api.list_namespaced_pod( - namespace=self.namespace, + namespace=self.definition.namespace, label_selector=f"job-name={self.external_id}", ) except k8s_client.ApiException: @@ -282,16 +277,15 @@ def collect_output(self) -> ProcessOutput | None: """Return the combined Kubernetes container log for this Job.""" if not self.external_id: return None - _, core_api = self.runtime.get_clients() - pods = core_api.list_namespaced_pod( - namespace=self.namespace, + 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 = core_api.read_namespaced_pod_log( + response = self.runtime.core_api.read_namespaced_pod_log( name=pods.items[0].metadata.name, - namespace=self.namespace, + namespace=self.definition.namespace, container="task", _preload_content=False, ) @@ -302,11 +296,10 @@ def stop(self) -> None: """Delete this Kubernetes Job and its associated Pods.""" if not self.external_id: return - batch_api, _ = self.runtime.get_clients() try: - batch_api.delete_namespaced_job( + self.runtime.batch_api.delete_namespaced_job( name=self.external_id, - namespace=self.namespace, + namespace=self.definition.namespace, body=k8s_client.V1DeleteOptions(propagation_policy="Background"), ) except k8s_client.ApiException as error: @@ -333,6 +326,23 @@ class KubernetesRuntimeBundle(BaseRuntimeBundle): 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 diff --git a/src/beakerhub/services/dashboard/aws_ecs_dashboard.py b/src/beakerhub/services/dashboard/aws_ecs_dashboard.py index 2ad72b4..195df44 100644 --- a/src/beakerhub/services/dashboard/aws_ecs_dashboard.py +++ b/src/beakerhub/services/dashboard/aws_ecs_dashboard.py @@ -1,15 +1,115 @@ -"""AWS ECS implementation skeleton for the dashboard service.""" +"""AWS ECS implementation of the dashboard service.""" +from datetime import datetime, timezone from typing import Any -from beakerhub.services.dashboard.base import BaseDashboardService +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 to the BeakerHub dashboard.""" + """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]: - raise NotImplementedError("AWS ECS dashboard data is not implemented") + 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, @@ -17,4 +117,206 @@ def get_session_logs( container: str, tail_lines: int, ) -> dict[str, Any]: - raise NotImplementedError("AWS ECS session logs are not implemented") + 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 index e7d695f..dcdd818 100644 --- a/src/beakerhub/services/dashboard/base.py +++ b/src/beakerhub/services/dashboard/base.py @@ -18,7 +18,12 @@ class BaseDashboardService(LoggingConfigurable): """Provide runtime inventory and session logs to dashboard handlers.""" def get_dashboard(self) -> dict[str, Any]: - """Return provider-specific dashboard data.""" + """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( diff --git a/src/beakerhub/services/dashboard/handlers.py b/src/beakerhub/services/dashboard/handlers.py index df1a38d..800436a 100644 --- a/src/beakerhub/services/dashboard/handlers.py +++ b/src/beakerhub/services/dashboard/handlers.py @@ -147,6 +147,7 @@ def compute_etag(self) -> None: 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)} @@ -154,6 +155,40 @@ async def get(self): 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.""" @@ -169,9 +204,19 @@ async def get(self, owner: str, session_id: str): 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( - session_id, + runtime_id, container, tail_lines, ) @@ -193,6 +238,11 @@ async def get(self, owner: str, session_id: str): 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, diff --git a/src/beakerhub/services/dashboard/kubernetes_dashboard.py b/src/beakerhub/services/dashboard/kubernetes_dashboard.py index eb121d8..d08a057 100644 --- a/src/beakerhub/services/dashboard/kubernetes_dashboard.py +++ b/src/beakerhub/services/dashboard/kubernetes_dashboard.py @@ -5,6 +5,7 @@ from traitlets import Unicode +from beakerhub.runtimes.kubernetes import KubernetesRuntime from beakerhub.services.dashboard.base import ( BaseDashboardService, DashboardServiceError, @@ -22,20 +23,19 @@ class KubernetesDashboardService(BaseDashboardService): help="Kubernetes namespace inspected by the dashboard service.", ) + def _runtime(self) -> KubernetesRuntime: + parent_runtime = getattr(self.parent, "runtime", None) + if isinstance(parent_runtime, KubernetesRuntime): + return parent_runtime + return KubernetesRuntime(namespace=self.namespace) + def get_dashboard(self) -> dict[str, Any]: - from kubernetes import client as k8s_client 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"} namespace = self.namespace @@ -54,20 +54,85 @@ def get_dashboard(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) @@ -389,17 +454,12 @@ def get_session_logs( 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 DashboardServiceError( - 503, - "No Kubernetes configuration found", - ) - - core_api = k8s_client.CoreV1Api() + raise DashboardServiceError( + 503, + "No Kubernetes configuration found", + ) namespace = self.namespace @@ -436,6 +496,7 @@ def get_session_logs( truncated = len(log_lines) >= tail_lines return { + "runtime_name": pod_name, "pod_name": pod_name, "container": container, "logs": logs or "", diff --git a/src/beakerhub/services/spawner/aws_ecs_spawner.py b/src/beakerhub/services/spawner/aws_ecs_spawner.py index 9067706..d47fe77 100644 --- a/src/beakerhub/services/spawner/aws_ecs_spawner.py +++ b/src/beakerhub/services/spawner/aws_ecs_spawner.py @@ -1,32 +1,44 @@ -"""ECS-backed JupyterHub spawner. +"""AWS ECS-backed JupyterHub spawner.""" -This module deliberately contains only ECS control-plane behavior. Moto can -exercise that behavior in unit tests; a Docker-backed ECS implementation such -as LocalStack is needed to verify that a task image actually starts. -""" - -from typing import Any +import asyncio +import hashlib +import os +import time +from functools import cache +from typing import Any, Literal import boto3 -from traitlets import Bool, Dict, List, Unicode +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 one ECS task. + """Launch a Beaker session as an ECS runtime process.""" - Networking/proxy registration is intentionally not implemented here yet. - ``start`` therefore launches and records the ECS task but does not claim - that the Jupyter server is reachable. - """ + runtime = Instance(AwsEcsRuntime, allow_none=False) - cluster_name = Unicode("default", config=True, help="ECS cluster name or ARN.") + cluster_name = Unicode(config=True, help="ECS cluster name or ARN.") task_definition = Unicode( config=True, - help="ECS task-definition family, revision, or ARN for notebook tasks.", + 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.", ) @@ -35,108 +47,504 @@ class BeakerAwsECSSpawner(BeakerhubImageSpawner): help="Tags applied to every ECS notebook task.", ) subnets = List( - Unicode(), default_value=[], config=True, + Unicode(), config=True, help="Subnets used by the task's awsvpc network configuration.", ) - security_groups = List(Unicode(), default_value=[], config=True) - assign_public_ip = Bool(False, config=True) - fargate = Bool(False, config=True) + 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) - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self.boto = boto3.client("ecs") - self.task_arn: str | None = None + @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) -> list[dict[str, str]]: + 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 [{"key": key, "value": value} for key, value in tags.items()] + return tags - def _network_configuration(self) -> dict[str, Any] | None: - if not self.subnets: - if self.fargate: - raise ValueError("BeakerAwsECSSpawner.subnets is required for Fargate tasks") - return None - return { - "awsvpcConfiguration": { - "subnets": self.subnets, - "securityGroups": self.security_groups, - "assignPublicIp": "ENABLED" if self.assign_public_ip else "DISABLED", + 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 _run_task_request(self) -> dict[str, Any]: - if not self.task_definition: - raise ValueError("BeakerAwsECSSpawner.task_definition must be configured") + 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") - request: dict[str, Any] = { - "cluster": self.cluster_name, - "taskDefinition": self.task_definition, - "count": 1, - "overrides": { - "containerOverrides": [{ - "name": self.container_name, - "environment": [ - {"name": name, "value": str(value)} - for name, value in self.get_env().items() - ], - }], + 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, }, - "tags": self._task_tags(), + 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}" + } } - if self.fargate: - request["launchType"] = "FARGATE" - network_configuration = self._network_configuration() - if network_configuration: - request["networkConfiguration"] = network_configuration - if self.volume_configs: - request["volumeConfigurations"] = self.volume_configs - return request - - def start(self) -> None: - """Request one ECS task and retain its ARN for later lifecycle calls.""" - response = self.boto.run_task(**self._run_task_request()) - tasks = response.get("tasks", []) - if not tasks: - failures = response.get("failures", []) - raise RuntimeError(f"ECS did not start a task: {failures!r}") - self.task_arn = tasks[0]["taskArn"] - - def stop(self, now: bool = False) -> None: - """Stop the task, if this spawner has launched one.""" + 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: - self.boto.stop_task( - cluster=self.cluster_name, - task=self.task_arn, - reason="Stopped by BeakerHub server", - ) + process = self._process() + process.stop() + await process.await_completion() - def poll(self) -> int | None: - """Return ``None`` while ECS is provisioning/running, otherwise an exit code.""" + async def poll(self) -> int | None: + """Map normalized ECS process status to JupyterHub poll semantics.""" if not self.task_arn: return 0 - response = self.boto.describe_tasks(cluster=self.cluster_name, tasks=[self.task_arn]) - task = next( - (item for item in response.get("tasks", []) if item.get("taskArn") == self.task_arn), - None, - ) - if task is None: - return 1 - if task.get("lastStatus") in {"PROVISIONING", "PENDING", "ACTIVATING", "RUNNING"}: + status: ProcessStatus = self._process().status + if status.state in {"pending", "running"}: return None - for container in task.get("containers", []): - exit_code = container.get("exitCode") - if exit_code is not None: - return int(exit_code) - return 1 + 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() @@ -151,3 +559,55 @@ def load_state(self, state: dict[str, Any]) -> None: 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/services/spawner/base.py b/src/beakerhub/services/spawner/base.py index 624e9da..33c0a29 100644 --- a/src/beakerhub/services/spawner/base.py +++ b/src/beakerhub/services/spawner/base.py @@ -4,6 +4,7 @@ from jupyterhub.spawner import Spawner from traitlets import Dict, Unicode, default, validate +from beakerhub import orm from beakerhub.auth.user import BeakerhubUser from beakerhub.services.secrets import VAULT_ENV_VAR_LIST_KEY @@ -80,7 +81,6 @@ def _extend_env(self, env: dict[str, str]) -> dict[str, str]: return env - def start(self): raise NotImplementedError() @@ -90,6 +90,81 @@ def stop(self, now=False): 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 = 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 + 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 = ( + self.db.query(orm.Context) + .filter(orm.Context.slug == context_slug) + .first() + ) + if context_record: + 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, + ) + ).fetchall() + disabled_secret_ids = {row.node_secret_id for row in disabled_rows} + + secrets_env: dict[str, str] = {} + # Policy overrides travel separately from the values: the value goes into the pod + # environment, while the policies tell the node what it may do with that value. + # Both are keyed by env_var, so a node-specific secret overrides a global one in + # exactly the same way for each. + secrets_policies: dict[str, dict[str, str]] = {} + + def collect(secret: orm.NodeSecret) -> None: + if secret.id in disabled_secret_ids: + return + secrets_env[secret.env_var] = secret.value + # An empty dict means "every axis at its default", which the node already + # assumes, so there is nothing to send. + if secret.policies: + secrets_policies[secret.env_var] = secret.policies + else: + # A node-specific secret with no overrides must not inherit the policies + # of the global secret it shadows. + secrets_policies.pop(secret.env_var, None) + + # Global secrets (node_image_id IS NULL) + global_secrets = ( + self.db.query(orm.NodeSecret) + .filter(orm.NodeSecret.node_image_id.is_(None)) + .all() + ) + for secret in global_secrets: + collect(secret) + + # Node-specific secrets (override globals) + if node_record: + node_secrets = ( + self.db.query(orm.NodeSecret) + .filter(orm.NodeSecret.node_image_id == node_record.id) + .all() + ) + for secret in node_secrets: + collect(secret) + + self.node_env = secrets_env or {} + self.node_policy_overrides = secrets_policies or {} class BeakerhubImageSpawner(BeakerSpawner): @@ -146,3 +221,15 @@ def _validate_image(self, proposal): 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 index 9b7b5db..45bad96 100644 --- a/src/beakerhub/services/spawner/kubernetes_spawner.py +++ b/src/beakerhub/services/spawner/kubernetes_spawner.py @@ -37,94 +37,17 @@ def node_type(self) -> str | None: return None return pod_name - @staticmethod - def apply_user_options(spawner: "BeakerKubeSpawner", 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 - - if (context := user_options.get("contextSlug", None)): - if ':' in context: - context = context.split(":")[-1] - context_slug = context - spawner.beaker_context = context - - # 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) - .filter(orm.Context.slug == context_slug) - .first() - ) - if context_record: - disabled_rows = spawner.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, - ) - ).fetchall() - disabled_secret_ids = {row.node_secret_id for row in disabled_rows} - - secrets_env: dict[str, str] = {} - # Policy overrides travel separately from the values: the value goes into the pod - # environment, while the policies tell the node what it may do with that value. - # Both are keyed by env_var, so a node-specific secret overrides a global one in - # exactly the same way for each. - secrets_policies: dict[str, dict[str, str]] = {} - - def collect(secret: orm.NodeSecret) -> None: - if secret.id in disabled_secret_ids: - return - secrets_env[secret.env_var] = secret.value - # An empty dict means "every axis at its default", which the node already - # assumes, so there is nothing to send. - if secret.policies: - secrets_policies[secret.env_var] = secret.policies - else: - # A node-specific secret with no overrides must not inherit the policies - # of the global secret it shadows. - secrets_policies.pop(secret.env_var, None) - - # Global secrets (node_image_id IS NULL) - global_secrets = ( - spawner.db.query(orm.NodeSecret) - .filter(orm.NodeSecret.node_image_id.is_(None)) - .all() - ) - for secret in global_secrets: - collect(secret) - - # Node-specific secrets (override globals) - if node_record: - node_secrets = ( - spawner.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 {} + 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: + if (node_slug := user_options.get("nodeSlug", None)): secret_name = f"{node_slug}-secrets" - env_from: list = spawner.extra_container_config.setdefault("envFrom", []) + 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 } }) - - if (context_config := user_options.get("contextOptions", None)): - spawner.context_config = context_config diff --git a/src/beakerhub/services/task/aws_ecs_task_runner.py b/src/beakerhub/services/task/aws_ecs_task_runner.py index e0a7ced..e4fc9ec 100644 --- a/src/beakerhub/services/task/aws_ecs_task_runner.py +++ b/src/beakerhub/services/task/aws_ecs_task_runner.py @@ -1,388 +1,168 @@ -"""AWS ECS implementation skeleton for the task-runner service.""" +"""AWS ECS implementation of the task-runner service.""" -import typing -from typing import Any +from typing import Any, Literal -import boto3 import traitlets -from botocore.client import BaseClient -from botocore.exceptions import ClientError -from beakerhub.services.task.base import ( - BaseTaskRunnerService, - RunningTask, - TaskOutput, - TaskStatus, -) +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): - """Run BeakerHub background tasks as AWS ECS tasks.""" - - boto_client: BaseClient = traitlets.Instance(klass=BaseClient, config=False) - logs_client: BaseClient = traitlets.Instance(klass=BaseClient, config=False) - - log_group: str = traitlets.Unicode( - default_value="", - help="CloudWatch Logs group that receives ECS task container output.", - config=True, - ) - log_stream_prefix: str = traitlets.Unicode( - default_value="beakerhub-task", - help="CloudWatch Logs stream prefix for ECS task containers.", - config=True, - ) - execution_role_arn: str = traitlets.Unicode( - default_value="", - help=( - "IAM role ARN used by ECS to pull images and write container logs. " - "This role needs ECR access for private images and " - "logs:CreateLogStream and logs:PutLogEvents for the log group." - ), - config=True, - ) - task_role_arn: str = traitlets.Unicode( - default_value="", - help=( - "IAM role ARN assumed by the task container. Grant it access to " - "AWS services required by the task." - ), - config=True, - ) - subnets: list[str] = traitlets.List( - traitlets.Unicode(), - default_value=[], - help="Subnets used by Fargate task awsvpc network configuration.", - config=True, - ) - security_groups: list[str] = traitlets.List( - traitlets.Unicode(), - default_value=[], - help="Security groups used by Fargate task awsvpc network configuration.", - config=True, - ) - assign_public_ip: bool = traitlets.Bool( - False, - help="Assign a public IP address to Fargate tasks.", - config=True, - ) - cpu_architecture: typing.Literal["X86_64", "ARM64"] = traitlets.Enum( - values=["X86_64", "ARM64"], - default_value="X86_64", - help="CPU architecture for task containers.", - config=True, - ) - task_definition_name: str = traitlets.Unicode( - default_value="beakerhub-tasks", - help="""Name of task definition "family". This is should not include the colon or revision number. Will be created if it does not exist.""", - config=True, - ) - cluster_name: str = traitlets.Unicode( - help="Name of ECS cluster to run task within", - config=True, - ) - task_group: str = traitlets.Unicode( - help="Optional. Name of task group to associate with the task.", - config=True, - ) - launch_type: typing.Literal["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"] = traitlets.Enum( - values=["EC2", "FARGATE", "EXTERNAL", "MANAGED_INSTANCES"], - default_value="FARGATE", - help="Method of launching ECS task", - config=True, - ) - launch_options: dict = traitlets.Dict( - help="Optional. If provided, overrides keyword arguments passed to boto3.client.run_task. See https://docs.aws.amazon.com/boto3/latest/reference/services/ecs/client/run_task.html", - config=True, - ) - task_overrides: dict = traitlets.Dict( - default_value=None, - allow_none=True, - help="Optional. ECS task TaskOverride values to be included. See https://docs.aws.amazon.com/AmazonECS/latest/APIReference/API_TaskOverride.html", - config=True, - ) - aws_region: str = traitlets.Unicode( - help="Optional. AWS Region cluster resides in.", - config=True, - ) - default_task_cpu: str = traitlets.Unicode( - default_value="4 vcpu", - help="Optional. Default CPU for task nodes. Used to create the task_definition if it does not exist. Does not override per-task runs.", - config=True, - ) - default_task_memory: str = traitlets.Unicode( - default_value="16GB", - help="Optional. Default memory for task nodes. Used to create the task_definition if it does not exist. Does not override per-task runs.", - config=True, - ) - default_task_ephemeral_storage: str = traitlets.Int( - default_value=50, - help="Optional. Default ephemeral storage for task nodes in GB. Used to create the task_definition if it does not exist. Does not override per-task runs.", - config=True, + """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("boto_client") - def _default_boto_client(self): - return boto3.client("ecs", region_name=self.aws_region or None) - - @traitlets.default("logs_client") - def _default_logs_client(self): - return boto3.client("logs", region_name=self.aws_region or None) + @traitlets.default("default_task_cpu") + def _default_task_cpu(self) -> str: + return self.runtime.default_task_cpu - @traitlets.default('launch_options') - def _default_launch_options(self): - return {} + @traitlets.default("default_task_memory") + def _default_task_memory(self) -> str: + return self.runtime.default_task_memory - @traitlets.default('aws_region') - def _default_aws_region(self): - session = boto3._get_default_session() - return (session.region_name if session else None) or "us-east-1" + @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) -> RunningTask: - """Register an image-specific definition and launch one ECS task.""" + 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 + ), + ) - self._validate_configuration() - definition_arn = self._register_task_definition(task) - request: dict[str, Any] = { - "cluster": self.cluster_name, - "taskDefinition": definition_arn, - "overrides": self._build_overrides(task), - "tags": [ - {"key": "beakerhub-task", "value": "true"}, - {"key": "beakerhub-task-type", "value": task.task_type}, - ], - } - if self.task_group: - request["group"] = self.task_group - if self.launch_type == "FARGATE" and "networkConfiguration" not in self.launch_options: - request["networkConfiguration"] = self._network_configuration() - if "capacityProviderStrategy" not in self.launch_options: - request["launchType"] = self.launch_type - request.update(self.launch_options) - - try: - response = self.boto_client.run_task(**request) - except ClientError as error: - raise RuntimeError(f"ECS failed to submit 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 RunningTask(tasks[0]["taskArn"], task, self) - - def _register_task_definition(self, task: BaseImageTask) -> str: - """Register the image selected by a task as an ECS definition revision.""" - container: dict[str, Any] = { - "name": "task-container", - "image": task.image, - "essential": True, - } - if task.entrypoint: - container["entryPoint"] = list(task.entrypoint) - if task.command: - container["command"] = list(task.command) - if task.working_directory: - container["workingDirectory"] = task.working_directory - if task.environment: - container["environment"] = [ - {"name": name, "value": value} - for name, value in task.environment.items() - ] - if self.log_group: - container["logConfiguration"] = { - "logDriver": "awslogs", - "options": { - "awslogs-group": self.log_group, - "awslogs-region": self.aws_region, - "awslogs-stream-prefix": self.log_stream_prefix, - }, - } - - request: dict[str, Any] = { - "family": self.task_definition_name, - "containerDefinitions": [container], - "cpu": self.default_task_cpu, - "memory": self.default_task_memory, - "runtimePlatform": { - "cpuArchitecture": self.cpu_architecture, - "operatingSystemFamily": "LINUX", - }, - } - if self.launch_type == "FARGATE": - request["ephemeralStorage"] = { - "sizeInGiB": self.default_task_ephemeral_storage - } - if self.execution_role_arn: - request["executionRoleArn"] = self.execution_role_arn - if self.task_role_arn: - request["taskRoleArn"] = self.task_role_arn - if self.launch_type == "FARGATE": - request.update({ - "networkMode": "awsvpc", - "requiresCompatibilities": ["FARGATE"], - }) - - try: - response = self.boto_client.register_task_definition(**request) - except ClientError as error: - raise RuntimeError(f"ECS failed to register task definition: {error}") from error - return response["taskDefinition"]["taskDefinitionArn"] - - def _build_overrides(self, task: BaseImageTask) -> dict[str, Any]: - overrides = dict(self.task_overrides or {}) - containers: list[dict[str, Any]] = [] - task_container: dict[str, Any] = {"name": "task-container"} - environment: dict[str, str] = {} - for container in overrides.pop("containerOverrides", []): - container = dict(container) - if container.get("name") != "task-container": - containers.append(container) - continue - environment.update({ - item["name"]: item["value"] - for item in container.pop("environment", []) - }) - task_container.update(container) - environment.update(task.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 _network_configuration(self) -> dict[str, Any]: - return { - "awsvpcConfiguration": { - "subnets": list(self.subnets), - "securityGroups": list(self.security_groups), - "assignPublicIp": "ENABLED" if self.assign_public_ip else "DISABLED", - } - } - - def _validate_configuration(self) -> None: - if not self.cluster_name.strip(): - raise ValueError("AwsEcsTaskRunnerService.cluster_name must be configured") - if not self.log_group: - raise ValueError("AwsEcsTaskRunnerService.log_group must be configured") - if not self.aws_region: - raise ValueError("AwsEcsTaskRunnerService.aws_region must be configured") - if not self.execution_role_arn: - raise ValueError( - "AwsEcsTaskRunnerService.execution_role_arn 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( - "AwsEcsTaskRunnerService.subnets must be configured for Fargate tasks" - ) - elif self.default_task_ephemeral_storage: - raise ValueError( - "default_task_ephemeral_storage is supported only for Fargate tasks" - ) - if ( - self.launch_type == "MANAGED_INSTANCES" - and not self.launch_options.get("capacityProviderStrategy") - ): - raise ValueError( - "Managed Instances tasks require launch_options.capacityProviderStrategy" - ) - - def get_status(self, external_id: str) -> TaskStatus: - task = self._describe_task(external_id) - if task is None: - return TaskStatus("failed", f"ECS task {external_id} was not found") - status = task.get("lastStatus", "UNKNOWN") - if status in {"PROVISIONING", "PENDING", "ACTIVATING"}: - return TaskStatus("pending", f"ECS task is {status.lower()}") - if status in {"RUNNING", "DEACTIVATING", "DEPROVISIONING", "STOPPING"}: - return TaskStatus("running", f"ECS task is {status.lower()}") - if status != "STOPPED": - return TaskStatus("failed", f"ECS task has unexpected status {status!r}") - - container = self._task_container(task) - if container and container.get("exitCode") == 0: - return TaskStatus("completed", "ECS task completed successfully") - reason = (container or {}).get("reason") or task.get("stoppedReason") or "Unknown failure" - return TaskStatus("failed", f"ECS task failed ({reason})") - - def get_output(self, external_id: str) -> TaskOutput | None: - """Return task output from the configured CloudWatch Logs group.""" - if not self.log_group: - return None - task = self._describe_task(external_id) - if task is None: - return None - container = self._task_container(task) - task_id = external_id.split('/')[-1] - container_name = (container or {}).get('name', 'task-container') - stream_name = f"{self.log_stream_prefix}/{container_name}/{task_id}" - - messages: list[str] = [] - next_token: str | None = None - while True: - request: dict[str, Any] = { - "logGroupName": self.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 for ECS task {external_id}: {error}" - ) from error - messages.extend(event["message"] for event in response.get("events", [])) - token = response.get("nextForwardToken") - if not token or token == next_token: - break - next_token = token - return TaskOutput(stdout="\n".join(messages), stderr="") - - def delete(self, external_id: str) -> None: - try: - self.boto_client.stop_task( - cluster=self.cluster_name, - task=external_id, - reason="Completed BeakerHub task cleanup", - ) - except ClientError as error: - raise RuntimeError(f"ECS failed to stop task {external_id}: {error}") from error - - def reap_stale_tasks(self) -> None: - """Stale-task reconciliation needs persisted task ownership and is deferred.""" - self.log.warning("ECS stale-task reaping is not implemented") - # TODO: Clean stale task-definition revisions as part of reaping tasks - - def _describe_task(self, external_id: str) -> dict[str, Any] | None: - try: - response = self.boto_client.describe_tasks( - cluster=self.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 - - @staticmethod - def _task_container(task: dict[str, Any]) -> dict[str, Any] | None: - return next( - (item for item in task.get("containers", []) if item.get("name") == "task-container"), - 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 index bcad92f..5aee026 100644 --- a/src/beakerhub/services/task/base.py +++ b/src/beakerhub/services/task/base.py @@ -1,95 +1,52 @@ -"""Base contracts for background task-runner services.""" +"""Task-runner adapters for persisted BeakerHub background tasks.""" -import asyncio -from dataclasses import dataclass -from time import monotonic -from typing import TYPE_CHECKING, Literal +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 -TaskState = Literal["pending", "running", "completed", "failed"] - - -@dataclass(frozen=True) -class TaskStatus: - """The runner's current view of a submitted task.""" - - state: TaskState - message: str | None = None - - @property - def done(self) -> bool: - """Return true when the task has reached a terminal state.""" - return self.state in {"completed", "failed"} - - -@dataclass(frozen=True) -class TaskOutput: - """Captured standard streams for a completed task.""" - - stdout: str - stderr: str - - -@dataclass(frozen=True) -class RunningTask: - """An in-memory handle for a task submitted to a runner.""" - - external_id: str - task_definition: "BaseTaskDefinition" - runner: "BaseTaskRunnerService" - - @property - def status(self) -> TaskStatus: - """Return the current status from the runner.""" - return self.runner.get_status(self.external_id) - - @property - def done(self) -> bool: - """Return true when the runner reports a terminal state.""" - return self.status.done - - @property - def output(self) -> TaskOutput | None: - """Return captured output after the runner reports task completion.""" - return self.runner.get_output(self.external_id) - - async def await_completion(self, timeout: float | None = 600) -> TaskStatus: - """Wait for terminal status, or raise ``TimeoutError``.""" - 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"Task {self.external_id!r} did not complete within {timeout} seconds" - ) - await asyncio.sleep(0.2) +# Keep task-service imports stable while using the runtime lifecycle contract. +TaskState: TypeAlias = ProcessState +TaskStatus = ProcessStatus +TaskOutput = ProcessOutput class BaseTaskRunnerService(LoggingConfigurable): - """Submit and manage background workloads for BeakerHub tasks.""" + """Adapt logical BeakerHub tasks to provider runtime processes. - def submit(self, task: "BaseTaskDefinition") -> RunningTask: - """Submit a task and return its in-memory runtime handle.""" - raise NotImplementedError + 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 get_status(self, external_id: str) -> TaskStatus: - """Return the current state and diagnostics for a submitted task.""" + def submit(self, task: "BaseTaskDefinition") -> BaseProcess: + """Launch a runtime process for a logical task.""" raise NotImplementedError - def get_output(self, external_id: str) -> TaskOutput | None: - """Return output for a terminal task, if the backend retains it.""" + 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: - """Remove a known submitted workload during normal task completion.""" - raise NotImplementedError + """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.""" diff --git a/src/beakerhub/services/task/handlers.py b/src/beakerhub/services/task/handlers.py index 6b35d7e..06f28f8 100644 --- a/src/beakerhub/services/task/handlers.py +++ b/src/beakerhub/services/task/handlers.py @@ -10,7 +10,7 @@ from tornado import web from beakerhub.orm import NodeImages, NodeImageTask -from beakerhub.tasks.image_import.task import ingest_import_output, launch_import_task +from beakerhub.tasks.image_import.task import ImageImportTask, launch_import_task log = logging.getLogger(__name__) @@ -66,11 +66,12 @@ async def get(self, image_id: str): 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") - task.result = ingest_import_output(self.db, node, output) + definition.on_success(task, status, output) task.status = "completed" task.error = None task.updated_at = datetime.now(timezone.utc) @@ -78,11 +79,8 @@ async def get(self, image_id: str): runner.delete(task.job_name) elif status.state == "failed": output = runner.get_output(task.job_name) - messages = [status.message] - if output and output.stderr: - messages.append(output.stderr) + definition.on_failure(task, status, output) task.status = "failed" - task.error = ". ".join(message for message in messages if message) task.updated_at = datetime.now(timezone.utc) self.db.commit() runner.delete(task.job_name) diff --git a/src/beakerhub/services/task/kubernetes_task_runner.py b/src/beakerhub/services/task/kubernetes_task_runner.py index 557284f..dc2d296 100644 --- a/src/beakerhub/services/task/kubernetes_task_runner.py +++ b/src/beakerhub/services/task/kubernetes_task_runner.py @@ -1,27 +1,24 @@ -"""Kubernetes implementation of the task-runner service.""" +"""Kubernetes Job implementation of the task-runner service.""" from typing import Any -from uuid import uuid4 -from kubernetes import client as k8s_client -from kubernetes import config as k8s_config -from traitlets import Dict, Integer, List, Unicode +from traitlets import Dict, Instance, Integer, List, Unicode, default -from beakerhub.services.task.base import ( - BaseTaskRunnerService, - RunningTask, - TaskOutput, - TaskStatus, +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): - """Run BeakerHub background tasks as Kubernetes Jobs.""" + """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 init container.", + help="Resource requests and limits for the node-image task container.", ) backoff_limit = Integer( 0, @@ -41,204 +38,65 @@ class KubernetesTaskRunnerService(BaseTaskRunnerService): node_selector = Dict(config=True, help="Node selector for task Pods.") tolerations: Any = List(config=True, help="Tolerations for task Pods.") namespace = Unicode( - "beakerhub", 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 - @staticmethod - def _get_clients() -> tuple[k8s_client.BatchV1Api, k8s_client.CoreV1Api]: - try: - k8s_config.load_incluster_config() - except k8s_config.ConfigException: - k8s_config.load_kube_config() - return k8s_client.BatchV1Api(), k8s_client.CoreV1Api() + @default("node_selector") + def _default_node_selector(self) -> dict[str, str]: + return dict(self.runtime.node_selector) - def submit(self, task: BaseTaskDefinition) -> RunningTask: + @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( - f"KubernetesTaskRunnerService only supports image tasks, not " + "KubernetesTaskRunnerService only supports image tasks, not " f"{task.task_type!r}" ) - external_id = self._submit_image_task(task) - return RunningTask( - external_id=external_id, - task_definition=task, - runner=self, - ) - - def _submit_image_task(self, task: BaseImageTask) -> str: - """Create a Job for a runtime-neutral image-task definition.""" - batch_api, _ = self._get_clients() - task_name = task.task_type.replace("_", "-") - job_name = f"beaker-task-{task_name}-{uuid4().hex[:8]}" - container = k8s_client.V1Container( - name="task", - image=task.image, - command=list(task.entrypoint) or None, - args=list(task.command) or None, - working_dir=task.working_directory, - env=[ - k8s_client.V1EnvVar(name=name, value=value) - for name, value in task.environment.items() - ] or None, - resources=self._build_resource_requirements( - dict(task.resources) or self.node_image_resources - ), - ) - pod_spec = k8s_client.V1PodSpec( - containers=[container], - restart_policy="Never", - node_selector=self.node_selector or None, - tolerations=( - [k8s_client.V1Toleration(**item) for item in self.tolerations] - if self.tolerations - else None - ), - ) - job = k8s_client.V1Job( - api_version="batch/v1", - kind="Job", - metadata=k8s_client.V1ObjectMeta( - name=job_name, - namespace=self.namespace, - labels={ - "app.kubernetes.io/name": "beakerhub", - "app.kubernetes.io/component": "task", - "beakerhub/task-type": task.task_type, - }, - ), - spec=k8s_client.V1JobSpec( - template=k8s_client.V1PodTemplateSpec( - metadata=k8s_client.V1ObjectMeta( - labels={ - "app.kubernetes.io/name": "beakerhub", - "app.kubernetes.io/component": "task", - } - ), - spec=pod_spec, - ), - backoff_limit=self.backoff_limit, - active_deadline_seconds=self.active_deadline_seconds, - ttl_seconds_after_finished=self.ttl_seconds_after_finished, - ), - ) - - batch_api.create_namespaced_job(namespace=self.namespace, body=job) - self.log.info( - "Created task job %s for task type %s", job_name, task.task_type - ) - return job_name - - def get_status(self, task_id: str) -> TaskStatus: - """Query the status of a Kubernetes Job.""" - batch_api, core_api = self._get_clients() - try: - job = batch_api.read_namespaced_job( - name=task_id, - namespace=self.namespace, - ) - except k8s_client.ApiException as error: - if error.status == 404: - return TaskStatus("failed", f"Job {task_id} not found") - raise - - status = job.status - if status.succeeded and status.succeeded > 0: - return TaskStatus("completed", "Job completed successfully") - if status.failed and status.failed > 0: - return TaskStatus("failed", self._get_failure_message(core_api, task_id)) - if status.active and status.active > 0: - return TaskStatus("running", "Job is running") - return TaskStatus("pending", "Job is pending") - - def get_output(self, task_id: str) -> TaskOutput | None: - """Return the task container log after its Job reaches a terminal state.""" - _, core_api = self._get_clients() - pods = core_api.list_namespaced_pod( + task_definition = KubernetesDefinition( + runtime=self.runtime, namespace=self.namespace, - label_selector=f"job-name={task_id}", - ) - if not pods.items: - return None - response = core_api.read_namespaced_pod_log( - name=pods.items[0].metadata.name, - namespace=self.namespace, - container="task", - _preload_content=False, + 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('_', '-')}", ) - logs = response.data.decode("utf-8") - # Kubernetes exposes a combined container log stream through this API. - return TaskOutput(stdout=logs or "", stderr="") + return KubernetesProcess.start(task_definition, runtime=self.runtime) - def _get_failure_message( - self, - core_api: k8s_client.CoreV1Api, - job_name: str, - ) -> str: - try: - pods = core_api.list_namespaced_pod( - namespace=self.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] - for status in pod.status.init_container_statuses or []: - message = self._container_failure_message(status, "Init container") - if message: - return message - for status in pod.status.container_statuses or []: - message = self._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() - - @staticmethod - def _container_failure_message(status, label: str) -> str | None: - 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 - - def delete(self, task_id: str) -> None: - """Delete a Job and its associated Pods.""" - batch_api, _ = self._get_clients() - try: - batch_api.delete_namespaced_job( - name=task_id, - namespace=self.namespace, - body=k8s_client.V1DeleteOptions(propagation_policy="Background"), - ) - except k8s_client.ApiException as error: - if error.status != 404: - raise - - @staticmethod - def _build_resource_requirements( - resources: dict, - ) -> k8s_client.V1ResourceRequirements | None: - if not resources: - return None - return k8s_client.V1ResourceRequirements( - requests=resources.get("requests"), - limits=resources.get("limits"), + 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/tasks/image_import/task.py b/src/beakerhub/tasks/image_import/task.py index 1964e9a..7e19262 100644 --- a/src/beakerhub/tasks/image_import/task.py +++ b/src/beakerhub/tasks/image_import/task.py @@ -6,10 +6,10 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, object_session from beakerhub.orm import NodeImages, NodeImageTask -from beakerhub.services.task.base import TaskOutput +from beakerhub.services.task.base import TaskOutput, TaskStatus from beakerhub.tasks.base import BaseImageTask from beakerhub.utils import ingest_interchange_dump @@ -22,14 +22,48 @@ @dataclass(frozen=True, kw_only=True) class ImageImportTask(BaseImageTask): - """Definition of a task that imports context metadata from a node image.""" + """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, @@ -61,14 +95,7 @@ def launch_import_task( db.commit() try: - running_task = app.task_runner.submit( - ImageImportTask( - image=node_image.default_img_string, - entrypoint=("sh", "-c"), - command=("beaker context dump",), - node_image=node_image, - ) - ) + 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}" diff --git a/tests/unit/runtimes/test_aws_ecs.py b/tests/unit/runtimes/test_aws_ecs.py index c48900d..ec4b54e 100644 --- a/tests/unit/runtimes/test_aws_ecs.py +++ b/tests/unit/runtimes/test_aws_ecs.py @@ -1,8 +1,10 @@ """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, @@ -37,30 +39,50 @@ def ecs_runtime(**overrides): return AwsEcsRuntime(**values) -def test_process_uses_runtime_defaults_and_allows_process_overrides(): +def test_definition_uses_runtime_defaults_and_allows_definition_overrides(): runtime = ecs_runtime() - definition = AwsEcsDefinition(task_definition="notebook:1") - - process = AwsEcsProcess(definition, runtime=runtime) - overridden = AwsEcsProcess( - definition, + 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="process-cluster", - subnets=["subnet-process"], + cluster_name="definition-cluster", + subnets=["subnet-definition"], ) - assert process.cluster_name == "runtime-cluster" - assert process.subnets == ["subnet-runtime"] - assert process.security_groups == ["sg-runtime"] - assert process.log_group == "/beakerhub/runtime" - assert process.launch_options == {"enableExecuteCommand": True} - assert overridden.cluster_name == "process-cluster" - assert overridden.subnets == ["subnet-process"] + 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"] - process.subnets.append("local-change") + 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( @@ -70,7 +92,8 @@ def test_fixed_definition_launch_uses_runtime_networking_and_process_metadata(): process = AwsEcsProcess.start( AwsEcsDefinition( - task_definition="notebook:1", + image="notebook:latest", + task_definition_arn="notebook:1", container_name="notebook", environment={"BEAKERHUB_USER": "ada"}, tags={"beaker-session": "session-1"}, @@ -103,16 +126,49 @@ def test_fixed_definition_launch_uses_runtime_networking_and_process_metadata(): } +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"}} ) - process = AwsEcsProcess( + runtime.ecs_client.run_task = Mock( + return_value={"tasks": [{"taskArn": "task-arn"}]} + ) + process = AwsEcsProcess.start( AwsEcsDefinition( image="example:latest", cpu="1 vcpu", @@ -122,7 +178,7 @@ def test_dynamic_definition_resources_override_runtime_defaults(): runtime=runtime, ) - assert process._register_task_definition() == "definition-arn" + assert process.external_id == "task-arn" request = runtime.ecs_client.register_task_definition.call_args.kwargs assert request["cpu"] == "1 vcpu" @@ -130,6 +186,88 @@ def test_dynamic_definition_resources_override_runtime_defaults(): 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( diff --git a/tests/unit/runtimes/test_base.py b/tests/unit/runtimes/test_base.py index 3086eaf..73940b6 100644 --- a/tests/unit/runtimes/test_base.py +++ b/tests/unit/runtimes/test_base.py @@ -1,6 +1,6 @@ """Tests for provider-neutral runtime contracts.""" -from dataclasses import dataclass +from traitlets import Unicode from beakerhub.runtimes.base import ( BaseDefinition, @@ -11,9 +11,8 @@ ) -@dataclass(frozen=True, kw_only=True) class ExampleDefinition(BaseDefinition): - value: str = "example" + value = Unicode("example", config=True) class ExampleRuntime(BaseRuntime): @@ -35,22 +34,30 @@ 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_shared_runtime(): +def test_bundle_creates_definition_and_process_with_application_runtime(): bundle = ExampleBundle() + runtime = ExampleRuntime() - definition = bundle.create_definition(value="configured") - process = bundle.start_process(definition) + 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.runtime is bundle.runtime + assert process.parent is runtime + assert process.runtime is runtime assert process.external_id == "example-process" assert process.status == ProcessStatus("running") @@ -60,3 +67,7 @@ def test_process_status_marks_only_terminal_states_done(): 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 index b746748..a683f28 100644 --- a/tests/unit/runtimes/test_kubernetes.py +++ b/tests/unit/runtimes/test_kubernetes.py @@ -70,7 +70,8 @@ def test_start_creates_job_and_records_the_provider_identifier(monkeypatch): batch_api = Mock() core_api = Mock() instance = runtime() - monkeypatch.setattr(instance, "get_clients", lambda: (batch_api, core_api)) + instance.batch_api = batch_api + instance.core_api = core_api process = KubernetesProcess.start( KubernetesDefinition(image="example:latest", name_prefix="image-import"), @@ -98,7 +99,8 @@ def test_describe_reports_job_completion(monkeypatch): status=SimpleNamespace(succeeded=1, failed=None, active=None) ) instance = runtime() - monkeypatch.setattr(instance, "get_clients", lambda: (batch_api, Mock())) + instance.batch_api = batch_api + instance.core_api = Mock() process = KubernetesProcess( KubernetesDefinition(image="example:latest"), runtime=instance, diff --git a/tests/unit/services/spawner/test_aws_ecs.py b/tests/unit/services/spawner/test_aws_ecs.py index 72da187..2677aa2 100644 --- a/tests/unit/services/spawner/test_aws_ecs.py +++ b/tests/unit/services/spawner/test_aws_ecs.py @@ -1,25 +1,53 @@ -"""Unit tests for the ECS control-plane spawner scaffold. - -These use botocore's request-validating Stubber rather than Docker: they are -fast tests for the ECS contract. A future LocalStack suite should cover actual -container execution separately. -""" +"""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 botocore.stub import Stubber +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 ECS-spawner methods.""" + """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", @@ -28,123 +56,508 @@ def configured_spawner(**overrides): "subnets": ["subnet-123"], "security_groups": ["sg-123"], "assign_public_ip": True, - "fargate": 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._network_configuration = lambda: BeakerAwsECSSpawner._network_configuration(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_run_task_request_contains_environment_tags_and_fargate_networking(): - spawner = configured_spawner() +def test_definition_contains_session_overrides_and_runtime_launch_options(): + spawner = configured_spawner( + volume_configs=[{"name": "session-volume", "managedEBSVolume": {}}] + ) - request = BeakerAwsECSSpawner._run_task_request(spawner) + definition = BeakerAwsECSSpawner._definition(spawner) - assert request == { - "cluster": "notebooks", - "taskDefinition": "beaker-notebook:3", - "count": 1, - "launchType": "FARGATE", - "networkConfiguration": { - "awsvpcConfiguration": { - "subnets": ["subnet-123"], - "securityGroups": ["sg-123"], - "assignPublicIp": "ENABLED", - } - }, - "overrides": { - "containerOverrides": [{ + 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", }, - "tags": [ - {"key": "environment", "value": "test"}, - {"key": "beaker-session", "value": "session-123"}, + 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) == [ - {"key": "beaker-session", "value": "session-123"} - ] + assert BeakerAwsECSSpawner._task_tags(spawner) == { + "beaker-session": "session-123" + } @pytest.mark.parametrize( - ("overrides", "message"), + "overrides", [ - ({"task_definition": ""}, "task_definition"), - ({"container_name": ""}, "container_name"), - ({"subnets": []}, "subnets"), + {"efs_volume_configuration": {"fileSystemId": "fs-123"}}, + {"efs_mount_path": "/home/beaker"}, + { + "efs_volume_configuration": {"rootDirectory": "/"}, + "efs_mount_path": "/home/beaker", + }, ], ) -def test_run_task_request_rejects_incomplete_fargate_configuration(overrides, message): +def test_definition_rejects_incomplete_efs_configuration(overrides): spawner = configured_spawner(**overrides) - with pytest.raises(ValueError, match=message): - BeakerAwsECSSpawner._run_task_request(spawner) + with pytest.raises(ValueError, match="efs_"): + BeakerAwsECSSpawner._definition(spawner) -def test_start_stop_and_poll_use_ecs_control_plane(): - client = boto3.client( - "ecs", - region_name="us-east-1", - aws_access_key_id="test", - aws_secret_access_key="test", +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}]} ) - request = {"cluster": "notebooks", "taskDefinition": "beaker-notebook:3", "count": 1} - spawner = configured_spawner(boto=client, task_arn=None) - spawner._run_task_request = lambda: request + shared_runtime.ecs_client.register_task_definition = Mock() + spawner = configured_spawner(runtime=shared_runtime) - with Stubber(client) as stubber: - stubber.add_response("run_task", {"tasks": [{"taskArn": TASK_ARN}]}, request) - BeakerAwsECSSpawner.start(spawner) - assert spawner.task_arn == TASK_ARN + await BeakerAwsECSSpawner.start(spawner) - stubber.add_response( - "describe_tasks", - {"tasks": [{"taskArn": TASK_ARN, "lastStatus": "RUNNING"}]}, - {"cluster": "notebooks", "tasks": [TASK_ARN]}, - ) - assert BeakerAwsECSSpawner.poll(spawner) is None + 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() - stubber.add_response( - "stop_task", - {"task": {"taskArn": TASK_ARN, "lastStatus": "STOPPED"}}, - { - "cluster": "notebooks", - "task": TASK_ARN, - "reason": "Stopped by BeakerHub server", - }, - ) - BeakerAwsECSSpawner.stop(spawner) +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="") -def test_poll_returns_container_exit_code_for_stopped_task(): - client = boto3.client( - "ecs", region_name="us-east-1", aws_access_key_id="test", aws_secret_access_key="test" + await BeakerAwsECSSpawner.start(spawner) + + registered = shared_runtime.ecs_client.register_task_definition.call_args.kwargs + assert registered["family"] == ( + "beakerhub-notebook_" + "d25e4b5db3dca83c3aae6b9fe0d9c3e9fc019ced629567804f94092004ac0316" ) - spawner = configured_spawner(boto=client, task_arn=TASK_ARN) - with Stubber(client) as stubber: - stubber.add_response( - "describe_tasks", - { - "tasks": [{ - "taskArn": TASK_ARN, + 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", - "containers": [{"name": "notebook", "exitCode": 23}], - }] - }, - {"cluster": "notebooks", "tasks": [TASK_ARN]}, - ) - assert BeakerAwsECSSpawner.poll(spawner) == 23 + "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/services/spawner/test_kubernetes.py b/tests/unit/services/spawner/test_kubernetes.py index 8ce56ef..5b36f4a 100644 --- a/tests/unit/services/spawner/test_kubernetes.py +++ b/tests/unit/services/spawner/test_kubernetes.py @@ -10,6 +10,37 @@ 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: """Tests for BeakerKubeSpawner class.""" @@ -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 index 4b4590c..8241b89 100644 --- a/tests/unit/services/test_aws_ecs_task_runner.py +++ b/tests/unit/services/test_aws_ecs_task_runner.py @@ -1,11 +1,11 @@ -"""Tests for the ECS task-runner task-definition configuration.""" +"""Tests for the ECS task-runner runtime adapter.""" from dataclasses import dataclass from unittest.mock import Mock import boto3 -import pytest +from beakerhub.runtimes.aws_ecs import AwsEcsRuntime from beakerhub.services.task.aws_ecs_task_runner import AwsEcsTaskRunnerService from beakerhub.tasks.base import BaseImageTask @@ -17,154 +17,111 @@ def task_type(self) -> str: return "example" -def test_task_definition_includes_configured_iam_roles(): +def runtime(**overrides) -> AwsEcsRuntime: ecs = boto3.client( "ecs", region_name="us-east-1", aws_access_key_id="test", aws_secret_access_key="test", ) - ecs.register_task_definition = Mock( - return_value={ - "taskDefinition": {"taskDefinitionArn": "task-definition-arn"} - } - ) - runner = AwsEcsTaskRunnerService( - boto_client=ecs, - log_group="/beakerhub/tasks", - execution_role_arn="arn:aws:iam::123456789012:role/ecs-task-execution", - task_role_arn="arn:aws:iam::123456789012:role/beakerhub-task", - ) - - assert runner._register_task_definition(ExampleImageTask(image="example:latest")) == ( - "task-definition-arn" - ) - - request = ecs.register_task_definition.call_args.kwargs - assert request["executionRoleArn"] == ( - "arn:aws:iam::123456789012:role/ecs-task-execution" + logs = boto3.client( + "logs", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", ) - assert request["taskRoleArn"] == "arn:aws:iam::123456789012:role/beakerhub-task" - - -def configured_runner(**overrides): values = { - "cluster_name": "tasks", - "log_group": "/beakerhub/tasks", - "aws_region": "us-east-1", - "execution_role_arn": "arn:aws:iam::123456789012:role/ecs-task-execution", - "subnets": ["subnet-123"], + "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 AwsEcsTaskRunnerService(**values) + return AwsEcsRuntime(**values) -def test_submit_uses_configured_fargate_networking(): - ecs = boto3.client( - "ecs", - region_name="us-east-1", - aws_access_key_id="test", - aws_secret_access_key="test", - ) - ecs.run_task = Mock(return_value={"tasks": [{"taskArn": "task-arn"}]}) - runner = configured_runner( - boto_client=ecs, - security_groups=["sg-123"], - assign_public_ip=True, - ) - runner._register_task_definition = Mock(return_value="task-definition-arn") +def test_runner_defaults_to_the_shared_runtime_configuration(): + shared_runtime = runtime() + runner = AwsEcsTaskRunnerService(runtime=shared_runtime) - runner.submit(ExampleImageTask(image="example:latest")) - - request = ecs.run_task.call_args.kwargs - assert request["networkConfiguration"] == { - "awsvpcConfiguration": { - "subnets": ["subnet-123"], - "securityGroups": ["sg-123"], - "assignPublicIp": "ENABLED", - } - } + 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_task_definition_uses_configured_cpu_architecture(): - ecs = boto3.client( - "ecs", - region_name="us-east-1", - aws_access_key_id="test", - aws_secret_access_key="test", +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": []} ) - ecs.register_task_definition = Mock( - return_value={"taskDefinition": {"taskDefinitionArn": "task-definition-arn"}} + shared_runtime.ecs_client.register_task_definition = Mock( + return_value={"taskDefinition": {"taskDefinitionArn": "definition-arn"}} ) - runner = configured_runner(boto_client=ecs, cpu_architecture="ARM64") - - runner._register_task_definition(ExampleImageTask(image="example:latest")) - - request = ecs.register_task_definition.call_args.kwargs - assert request["runtimePlatform"]["cpuArchitecture"] == "ARM64" - - -def test_task_overrides_merge_the_task_container_environment(): - runner = configured_runner( - task_overrides={ - "containerOverrides": [ - { - "name": "task-container", - "command": ["configured-command"], - "environment": [ - {"name": "FROM_CONFIG", "value": "configured"}, - {"name": "SHARED", "value": "configured"}, - ], - }, - {"name": "sidecar", "command": ["sidecar-command"]}, - ] - } + 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, ) - overrides = runner._build_overrides( + process = runner.submit( ExampleImageTask( image="example:latest", - environment={"SHARED": "task", "FROM_TASK": "task"}, + environment={"FROM_TASK": "value"}, ) ) - assert overrides["containerOverrides"] == [ - {"name": "sidecar", "command": ["sidecar-command"]}, + 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", - "command": ["configured-command"], - "environment": [ - {"name": "FROM_CONFIG", "value": "configured"}, - {"name": "SHARED", "value": "task"}, - {"name": "FROM_TASK", "value": "task"}, - ], - }, + "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") -@pytest.mark.parametrize( - ("overrides", "message"), - [ - ({"cluster_name": ""}, "cluster_name"), - ({"log_group": ""}, "log_group"), - ({"aws_region": ""}, "aws_region"), - ({"execution_role_arn": ""}, "execution_role_arn"), - ({"subnets": []}, "subnets"), - ( - {"launch_type": "EC2"}, - "default_task_ephemeral_storage", - ), - ( - { - "launch_type": "MANAGED_INSTANCES", - "default_task_ephemeral_storage": 0, - }, - "capacityProviderStrategy", - ), - ], -) -def test_configuration_validation_rejects_incompatible_settings(overrides, message): - runner = configured_runner(**overrides) - - with pytest.raises(ValueError, match=message): - runner._validate_configuration() + 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 index 204da91..e978b94 100644 --- a/tests/unit/services/test_service_wiring.py +++ b/tests/unit/services/test_service_wiring.py @@ -5,7 +5,13 @@ 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 @@ -38,6 +44,75 @@ def test_application_registers_service_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 diff --git a/tests/unit/services/test_task_base.py b/tests/unit/services/test_task_base.py index fb42da0..1b1d357 100644 --- a/tests/unit/services/test_task_base.py +++ b/tests/unit/services/test_task_base.py @@ -1,57 +1,78 @@ -"""Tests for task-runner lifecycle value objects.""" - -from dataclasses import dataclass +"""Tests for task-runner delegation to runtime processes.""" import pytest -from beakerhub.services.task.base import ( - BaseTaskRunnerService, - RunningTask, - TaskOutput, - TaskStatus, +from beakerhub.runtimes.base import ( + BaseDefinition, + BaseProcess, + BaseRuntime, + ProcessOutput, + ProcessStatus, ) -from beakerhub.tasks.base import BaseTaskDefinition +from beakerhub.services.task.base import BaseTaskRunnerService, TaskOutput, TaskStatus -@dataclass(frozen=True) -class ExampleTask(BaseTaskDefinition): - @property - def task_type(self) -> str: - return "example" +class CompleteProcess(BaseProcess): + def describe(self) -> ProcessStatus: + return ProcessStatus("completed") + def collect_output(self) -> ProcessOutput: + return ProcessOutput("output", "") -class CompleteTaskRunner(BaseTaskRunnerService): - def get_status(self, external_id: str) -> TaskStatus: - return TaskStatus("completed") + def stop(self) -> None: + return None - def get_output(self, external_id: str) -> TaskOutput: - return TaskOutput("output", "") +class PendingProcess(CompleteProcess): + def describe(self) -> ProcessStatus: + return ProcessStatus("running") -def test_running_task_uses_its_external_id_for_runner_operations(): - runner = CompleteTaskRunner() - task = RunningTask("runtime-123", ExampleTask(), runner) - assert task.status == TaskStatus("completed") - assert task.done is True - assert task.output == TaskOutput("output", "") +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 -@pytest.mark.asyncio -async def test_running_task_returns_terminal_status(): - runner = CompleteTaskRunner() - task = RunningTask("runtime-123", ExampleTask(), runner) - assert await task.await_completion() == TaskStatus("completed") +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_running_task_raises_on_timeout(): - class PendingTaskRunner(CompleteTaskRunner): - def get_status(self, external_id: str) -> TaskStatus: - return TaskStatus("running") +async def test_process_returns_terminal_status(): + process = complete_process() + + assert await process.await_completion() == ProcessStatus("completed") - task = RunningTask("runtime-123", ExampleTask(), PendingTaskRunner()) + +@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 task.await_completion(timeout=0) + await process.await_completion(timeout=0) 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'; - } -}