Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 33 additions & 15 deletions server/src/services/k8s/agent_sandbox_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
ApiException,
)

from src.config import AppConfig, IngressConfig, ExecdInitResources
from src.config import AppConfig
from src.services.helpers import format_ingress_endpoint
from src.api.schema import Endpoint, ImageSpec, NetworkPolicy
from src.services.k8s.agent_sandbox_template import AgentSandboxTemplateManager
Expand Down Expand Up @@ -82,17 +82,16 @@ class AgentSandboxProvider(WorkloadProvider):
def __init__(
self,
k8s_client: K8sClient,
template_file_path: Optional[str] = None,
shutdown_policy: str = "Delete",
service_account: Optional[str] = None,
ingress_config: Optional[IngressConfig] = None,
enable_informer: bool = True,
informer_factory: Optional[Callable[[str], WorkloadInformer]] = None,
informer_resync_seconds: int = 300,
informer_watch_timeout_seconds: int = 60,
app_config: Optional[AppConfig] = None,
execd_init_resources: Optional[ExecdInitResources] = None,
*,
informer_factory: Optional[Callable[[str], WorkloadInformer]] = None,
):
"""
Initialize AgentSandbox provider.

Configuration is read from app_config (kubernetes.*, agent_sandbox, ingress).
No separate config objects are passed.
"""
self.k8s_client = k8s_client
self.custom_api = k8s_client.get_custom_objects_api()
self.core_api = k8s_client.get_core_v1_api()
Expand All @@ -101,12 +100,31 @@ def __init__(
self.version = "v1alpha1"
self.plural = "sandboxes"

self.shutdown_policy = shutdown_policy
self.service_account = service_account
k8s_config = app_config.kubernetes if app_config else None
agent_config = app_config.agent_sandbox if app_config else None
self.shutdown_policy = (
agent_config.shutdown_policy if agent_config else "Delete"
)
self.service_account = (
k8s_config.service_account if k8s_config else None
)
template_file_path = (
agent_config.template_file if agent_config else None
)
self.template_manager = AgentSandboxTemplateManager(template_file_path)
self.ingress_config = ingress_config
self.execd_init_resources = execd_init_resources
self._enable_informer = enable_informer
self.ingress_config = app_config.ingress if app_config else None
self.execd_init_resources = (
k8s_config.execd_init_resources if k8s_config else None
)
self._enable_informer = (
k8s_config.informer_enabled if k8s_config else True
)
informer_resync_seconds = (
k8s_config.informer_resync_seconds if k8s_config else 300
)
informer_watch_timeout_seconds = (
k8s_config.informer_watch_timeout_seconds if k8s_config else 60
)
self._informer_factory = informer_factory or (
lambda ns: WorkloadInformer(
custom_api=self.custom_api,
Expand Down
46 changes: 30 additions & 16 deletions server/src/services/k8s/batchsandbox_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
ApiException,
)

from src.config import AppConfig, IngressConfig, INGRESS_MODE_GATEWAY, ExecdInitResources
from src.config import AppConfig, INGRESS_MODE_GATEWAY
from src.services.helpers import format_ingress_endpoint
from src.api.schema import Endpoint, ImageSpec, NetworkPolicy
from src.services.k8s.batchsandbox_template import BatchSandboxTemplateManager
Expand Down Expand Up @@ -60,26 +60,28 @@ class BatchSandboxProvider(WorkloadProvider):
def __init__(
self,
k8s_client: K8sClient,
template_file_path: Optional[str] = None,
ingress_config: Optional[IngressConfig] = None,
enable_informer: bool = True,
informer_factory: Optional[Callable[[str], WorkloadInformer]] = None,
informer_resync_seconds: int = 300,
informer_watch_timeout_seconds: int = 60,
app_config: Optional[AppConfig] = None,
execd_init_resources: Optional[ExecdInitResources] = None,
*,
informer_factory: Optional[Callable[[str], WorkloadInformer]] = None,
):
"""
Initialize BatchSandbox provider.

Configuration is read from app_config (kubernetes.*, ingress). No separate
config objects are passed.

Args:
k8s_client: Kubernetes client wrapper
template_file_path: Optional path to BatchSandbox CR YAML template file
app_config: Optional application config for secure runtime
k8s_client: Kubernetes client wrapper.
app_config: Application config; kubernetes and ingress settings are read from it.
informer_factory: Optional custom informer factory (for tests).
"""
self.k8s_client = k8s_client
self.custom_api = k8s_client.get_custom_objects_api()
self.ingress_config = ingress_config
k8s_config = app_config.kubernetes if app_config else None
self.ingress_config = app_config.ingress if app_config else None
self.execd_init_resources = (
k8s_config.execd_init_resources if k8s_config else None
)

# Initialize secure runtime resolver
self.resolver = SecureRuntimeResolver(app_config) if app_config else None
Expand All @@ -91,11 +93,23 @@ def __init__(
self.group = "sandbox.opensandbox.io"
self.version = "v1alpha1"
self.plural = "batchsandboxes"

# Template manager

template_file_path = (
k8s_config.batchsandbox_template_file if k8s_config else None
)
if template_file_path:
logger.info("Using BatchSandbox template file: %s", template_file_path)
self.template_manager = BatchSandboxTemplateManager(template_file_path)
self.execd_init_resources = execd_init_resources
self._enable_informer = enable_informer

self._enable_informer = (
k8s_config.informer_enabled if k8s_config else True
)
informer_resync_seconds = (
k8s_config.informer_resync_seconds if k8s_config else 300
)
informer_watch_timeout_seconds = (
k8s_config.informer_watch_timeout_seconds if k8s_config else 60
)
self._informer_factory = informer_factory or (
lambda ns: WorkloadInformer(
custom_api=self.custom_api,
Expand Down
3 changes: 0 additions & 3 deletions server/src/services/k8s/kubernetes_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,9 +110,6 @@ def __init__(self, config: Optional[AppConfig] = None):
self.workload_provider = create_workload_provider(
provider_type=provider_type,
k8s_client=self.k8s_client,
k8s_config=self.app_config.kubernetes,
agent_sandbox_config=self.app_config.agent_sandbox,
ingress_config=self.ingress_config,
app_config=self.app_config,
)
logger.info(
Expand Down
61 changes: 14 additions & 47 deletions server/src/services/k8s/provider_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import logging
from typing import Dict, Type, Optional

from src.config import AppConfig, KubernetesRuntimeConfig, AgentSandboxRuntimeConfig, IngressConfig
from src.config import AppConfig
from src.services.k8s.workload_provider import WorkloadProvider
from src.services.k8s.batchsandbox_provider import BatchSandboxProvider
from src.services.k8s.agent_sandbox_provider import AgentSandboxProvider
Expand All @@ -43,27 +43,26 @@
def create_workload_provider(
provider_type: str | None,
k8s_client: K8sClient,
k8s_config: Optional[KubernetesRuntimeConfig] = None,
agent_sandbox_config: Optional[AgentSandboxRuntimeConfig] = None,
ingress_config: Optional[IngressConfig] = None,
app_config: Optional[AppConfig] = None,
) -> WorkloadProvider:
"""
Create a WorkloadProvider instance based on the provider type.

All provider-specific configuration is read from app_config; no separate
config objects are passed.

Args:
provider_type: Type of provider (e.g., 'batchsandbox', 'pod', 'job').
provider_type: Type of provider (e.g., 'batchsandbox', 'agent-sandbox').
If None, uses the first registered provider.
k8s_client: Kubernetes client instance
k8s_config: Optional Kubernetes runtime configuration (for template file paths, etc.)
agent_sandbox_config: Optional agent-sandbox configuration (for template/shutdown policy)
app_config: Optional application config for secure runtime
k8s_client: Kubernetes client instance.
app_config: Application config; used by built-in providers for
kubernetes, agent_sandbox, and ingress settings.

Returns:
WorkloadProvider instance
WorkloadProvider instance.

Raises:
ValueError: If provider_type is not supported or no providers are registered
ValueError: If provider_type is not supported or no providers are registered.
"""
# Use first registered provider if not specified
if provider_type is None:
Expand All @@ -87,43 +86,11 @@ def create_workload_provider(
provider_class = _PROVIDER_REGISTRY[provider_type_lower]
logger.info(f"Creating workload provider: {provider_class.__name__}")

# Special handling for BatchSandboxProvider - pass template file path
if provider_type_lower == PROVIDER_TYPE_BATCHSANDBOX:
template_file = k8s_config.batchsandbox_template_file if k8s_config else None
if template_file:
logger.info(f"Using BatchSandbox template file: {template_file}")
return provider_class(
k8s_client,
template_file_path=template_file,
ingress_config=ingress_config,
enable_informer=k8s_config.informer_enabled,
informer_resync_seconds=k8s_config.informer_resync_seconds,
informer_watch_timeout_seconds=k8s_config.informer_watch_timeout_seconds,
app_config=app_config,
execd_init_resources=k8s_config.execd_init_resources if k8s_config else None,
)

# Special handling for AgentSandboxProvider - pass agent-specific settings
if provider_type_lower == PROVIDER_TYPE_AGENT_SANDBOX:
agent_config = agent_sandbox_config or AgentSandboxRuntimeConfig()
return provider_class(
k8s_client,
template_file_path=agent_config.template_file,
shutdown_policy=agent_config.shutdown_policy,
service_account=k8s_config.service_account if k8s_config else None,
ingress_config=ingress_config,
enable_informer=k8s_config.informer_enabled if k8s_config else True,
informer_resync_seconds=(
k8s_config.informer_resync_seconds if k8s_config else 300
),
informer_watch_timeout_seconds=(
k8s_config.informer_watch_timeout_seconds if k8s_config else 60
),
app_config=app_config,
execd_init_resources=k8s_config.execd_init_resources if k8s_config else None,
)
# Built-in providers accept (k8s_client, app_config) and read config internally
if provider_type_lower in (PROVIDER_TYPE_BATCHSANDBOX, PROVIDER_TYPE_AGENT_SANDBOX):
return provider_class(k8s_client, app_config=app_config)

# Providers without ingress-specific needs
# Custom/other providers may only accept k8s_client
return provider_class(k8s_client)


Expand Down
78 changes: 77 additions & 1 deletion server/tests/k8s/fixtures/k8s_fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,12 @@
import pytest

from src.api.schema import CreateSandboxRequest, ImageSpec, ResourceLimits
from src.config import KubernetesRuntimeConfig
from src.config import (
AppConfig,
AgentSandboxRuntimeConfig,
KubernetesRuntimeConfig,
RuntimeConfig,
)
from src.services.k8s.client import K8sClient
from src.services.k8s.provider_factory import PROVIDER_TYPE_BATCHSANDBOX

Expand Down Expand Up @@ -62,6 +67,77 @@ def agent_sandbox_runtime_config():
)


def _make_app_config(
*,
kubernetes: KubernetesRuntimeConfig,
agent_sandbox: AgentSandboxRuntimeConfig | None = None,
) -> AppConfig:
"""Build AppConfig for K8s tests."""
return AppConfig(
runtime=RuntimeConfig(type="kubernetes", execd_image="test-execd:latest"),
kubernetes=kubernetes,
agent_sandbox=agent_sandbox,
)


@pytest.fixture
def app_config(k8s_runtime_config):
"""Application config for batchsandbox provider tests."""
return _make_app_config(kubernetes=k8s_runtime_config)


@pytest.fixture
def app_config_agent_sandbox(agent_sandbox_runtime_config, tmp_path):
"""Application config for agent-sandbox provider tests."""
template_file = tmp_path / "agent_sandbox_template.yaml"
template_file.write_text(
"""
metadata:
annotations:
managed-by: opensandbox
spec:
podTemplate:
spec:
nodeSelector:
workload: sandbox
"""
)
agent_config = AgentSandboxRuntimeConfig(
template_file=str(template_file),
shutdown_policy="Retain",
ingress_enabled=True,
)
return _make_app_config(
kubernetes=agent_sandbox_runtime_config,
agent_sandbox=agent_config,
)


@pytest.fixture
def app_config_with_batch_template(tmp_path):
"""Application config with BatchSandbox template file path set."""
template_file = tmp_path / "test_template.yaml"
template_file.write_text("""
apiVersion: execution.alibaba-inc.com/v1alpha1
kind: BatchSandbox
metadata:
name: test-template
spec:
template:
spec:
nodeSelector:
gpu: "true"
""")
k8s_config = KubernetesRuntimeConfig(
kubeconfig_path="/tmp/test-kubeconfig",
namespace="test-namespace",
service_account="test-sa",
workload_provider=PROVIDER_TYPE_BATCHSANDBOX,
batchsandbox_template_file=str(template_file),
)
return _make_app_config(kubernetes=k8s_config)


@pytest.fixture
def k8s_runtime_config_with_template(tmp_path):
"""Provide Kubernetes configuration with template file"""
Expand Down
Loading