added optional kube.infrastructure.context

caracal-pipeline · Sep 14, 2023 · c3fe211 · c3fe211
1 parent 7f8bcb1
commit c3fe211
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 11 deletions.
diff --git a/stimela/backends/kube/__init__.py b/stimela/backends/kube/__init__.py
@@ -43,11 +43,15 @@ def is_remote():
 
 def init(backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger):
     from . import infrastructure
-    infrastructure.init(backend, log)
+    global AVAILABLE, STATUS
+    if not infrastructure.init(backend, log):
+        AVAILABLE = False
+        STATUS = "initialization error"
 
 def close(backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger):
     from . import infrastructure
-    infrastructure.close(backend, log)
+    if not AVAILABLE:
+        infrastructure.close(backend, log)
 
 def cleanup(backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger):
     from . import infrastructure
@@ -61,13 +65,14 @@ def run(cab: 'stimela.kitchen.cab.Cab', params: Dict[str, Any], fqname: str,
 
 _kube_client = _kube_config = None
 
-def get_kube_api():
+def get_kube_api(context: Optional[str]=None):
     global _kube_client
     global _kube_config
+    global _kube_context
 
     if _kube_config is None:
         _kube_config = True
-        kubernetes.config.load_kube_config()
+        kubernetes.config.load_kube_config(context=context)
 
     return core_v1_api.CoreV1Api(), CustomObjectsApi()
 
@@ -120,8 +125,9 @@ class StartupOptions(object):
             report_pvcs: bool = True                  # report any transient PVCs
             cleanup_pvcs: bool = True                 # cleanup any transient PVCs
 
-        on_exit:    ExitOptions = ExitOptions()                     # startup behaviour options
-        on_startup: StartupOptions = StartupOptions()               # cleanup behaviour options
+        context:    Optional[str] = None                # k8s context -- use default if not given
+        on_exit:    ExitOptions = ExitOptions()         # startup behaviour options
+        on_startup: StartupOptions = StartupOptions()   # cleanup behaviour options
 
     @dataclass 
     class Volume(object):
@@ -174,13 +180,15 @@ class LocalMount(object):
 
 
     enable:         bool = True
+
+    # infrastructure settings are global and can't be changed per cab or per step
+    infrastructure: Infrastructure = Infrastructure()
+
     namespace:      Optional[str] = None
     dask_cluster:   Optional[DaskCluster] = None
     service_account: str = "compute-runner"
     kubectl_path:   str = "kubectl"
 
-    infrastructure: Infrastructure = Infrastructure()
-
     volumes:        Dict[str, Volume] = EmptyDictDefault()
 
     inject_files:   Dict[str, FileInjection] = EmptyDictDefault()

diff --git a/stimela/backends/kube/infrastructure.py b/stimela/backends/kube/infrastructure.py
@@ -15,6 +15,7 @@
 
 from kubernetes import client
 from kubernetes.client.rest import ApiException
+from kubernetes.config import ConfigException
 
 from .kube_utils import resolve_unit
 
@@ -49,18 +50,22 @@ def init(backend: StimelaBackendOptions, log: logging.Logger, cleanup: bool = Fa
     global klog
     klog = log.getChild("kube")
     kube = backend.kube
+    try:
+        kube_api, _ = run_kube.get_kube_api(kube.infrastructure.context) 
+    except ConfigException as exc:
+        log_exception(exc, log=klog)
+        log_exception(BackendError("error initializing kube backend", exc), log=klog)
+        return False
 
     if cleanup:
         klog.info("cleaning up backend")
     else:
-        atexit.register(close, kube, klog)
+        atexit.register(close, backend, klog)
         klog.info("initializing kube backend")
 
     if cleanup or kube.infrastructure.on_startup.report_pods or kube.infrastructure.on_startup.cleanup_pods:
         klog.info("checking for k8s pods from other sessions")
 
-        kube_api, _ = run_kube.get_kube_api() 
-
         try:
             pods = kube_api.list_namespaced_pod(namespace=kube.namespace, 
                                                 label_selector=f"stimela_user={session_user}")
@@ -110,6 +115,8 @@ def init(backend: StimelaBackendOptions, log: logging.Logger, cleanup: bool = Fa
     if not cleanup and kube.volumes:
         resolve_volumes(kube, log=klog, refresh=False) # no refresh needed
 
+    return True
+
 
 def refresh_pvc_list(kube: KubeBackendOptions):
     kube_api, _ = get_kube_api()