polyaxon 2.0.0rc49__py3-none-any.whl → 2.4.0rc1__py3-none-any.whl

polyaxon/_k8s/manager/async_manager.py

@@ -38,7 +38,7 @@ class AsyncK8sManager(BaseK8sManager):
         if self.api_client:
             await self.api_client.close()
 
-    async def get_version(self, reraise=False):
+    async def get_version(self, reraise: bool = False):
         try:
             version = await self.k8s_version_api.get_code()
             return version.to_dict()
@@ -47,48 +47,68 @@ class AsyncK8sManager(BaseK8sManager):
             if reraise:
                 raise PolyaxonK8sError("Connection error: %s" % e) from e
 
-    async def get_pod(self, name, reraise=False) -> Optional[client.V1Pod]:
+    async def get_pod(
+        self, name, reraise: bool = False, namespace: str = None
+    ) -> Optional[client.V1Pod]:
         try:
             return await self.k8s_api.read_namespaced_pod(  # type: ignore[attr-defined]
-                name=name, namespace=self.namespace
+                name=name, namespace=namespace or self.namespace
             )
         except ApiException as e:
             if reraise:
                 raise PolyaxonK8sError("Connection error: %s" % e) from e
             return None
 
-    async def is_pod_running(self, pod_id: str, container_id: str) -> bool:
-        event = await self.k8s_api.read_namespaced_pod_status(pod_id, self.namespace)  # type: ignore[attr-defined]
+    async def is_pod_running(
+        self, pod_id: str, container_id: str, namespace: str = None
+    ) -> bool:
+        event = await self.k8s_api.read_namespaced_pod_status(pod_id, namespace=namespace or self.namespace)  # type: ignore[attr-defined]
         return is_pod_running(event, container_id)
 
     async def _list_namespace_resource(
-        self, resource_api, reraise=False, **kwargs
+        self, resource_api, reraise: bool = False, namespace: str = None, **kwargs
     ) -> List:
         try:
-            res = await resource_api(namespace=self.namespace, **kwargs)
-            return [p for p in res.items]
+            res = await resource_api(namespace=namespace or self.namespace, **kwargs)
+            if isinstance(res, dict):
+                items = res["items"]
+            else:
+                items = res.items
+            return [p for p in items]
         except ApiException as e:
             logger.error("K8S error: {}".format(e))
             if reraise:
                 raise PolyaxonK8sError("Connection error: %s" % e) from e
             return []
 
-    async def list_pods(self, reraise=False, **kwargs) -> List[client.V1Pod]:
+    async def list_pods(
+        self, reraise: bool = False, namespace: str = None, **kwargs
+    ) -> List[client.V1Pod]:
         return await self._list_namespace_resource(
             resource_api=self.k8s_api.list_namespaced_pod,  # type: ignore[attr-defined]
             reraise=reraise,
+            namespace=namespace,
             **kwargs,
         )
 
-    async def list_jobs(self, reraise=False, **kwargs) -> List[client.V1Job]:
+    async def list_jobs(
+        self, reraise: bool = False, namespace: str = None, **kwargs
+    ) -> List[client.V1Job]:
         return await self._list_namespace_resource(
             resource_api=self.k8s_batch_api.list_namespaced_job,  # type: ignore[attr-defined]
             reraise=reraise,
+            namespace=namespace,
             **kwargs,
         )
 
     async def list_custom_objects(
-        self, group, version, plural, reraise=False, **kwargs
+        self,
+        group,
+        version,
+        plural,
+        reraise: bool = False,
+        namespace: str = None,
+        **kwargs
     ) -> List:
         return await self._list_namespace_resource(
             resource_api=self.k8s_custom_object_api.list_namespaced_custom_object,  # type: ignore[attr-defined]
@@ -96,49 +116,64 @@ class AsyncK8sManager(BaseK8sManager):
             group=group,
             version=version,
             plural=plural,
+            namespace=namespace,
             **kwargs,
         )
 
     async def list_services(
-        self, reraise: bool = False, **kwargs
+        self, reraise: bool = False, namespace: str = None, **kwargs
     ) -> List[client.V1Service]:
         return await self._list_namespace_resource(
             resource_api=self.k8s_api.list_namespaced_service,  # type: ignore[attr-defined]
             reraise=reraise,
+            namespace=namespace,
             **kwargs,
         )
 
     async def list_deployments(
-        self, reraise: bool = False, **kwargs
+        self, reraise: bool = False, namespace: str = None, **kwargs
     ) -> List[client.V1Deployment]:
         return await self._list_namespace_resource(
             resource_api=self.k8s_apps_api.list_namespaced_deployment,  # type: ignore[attr-defined]
             reraise=reraise,
+            namespace=namespace,
             **kwargs,
         )
 
     async def create_custom_object(
-        self, name: str, group: str, version: str, plural: str, body: Dict
+        self,
+        name: str,
+        group: str,
+        version: str,
+        plural: str,
+        body: Dict,
+        namespace: str = None,
     ) -> Dict:
         resp = await self.k8s_custom_object_api.create_namespaced_custom_object(  # type: ignore[attr-defined]
             group=group,
             version=version,
             plural=plural,
-            namespace=self.namespace,
+            namespace=namespace or self.namespace,
             body=body,
         )
         logger.debug("Custom object `{}` was created".format(name))
         return resp
 
     async def update_custom_object(
-        self, name: str, group: str, version: str, plural: str, body: Dict
+        self,
+        name: str,
+        group: str,
+        version: str,
+        plural: str,
+        body: Dict,
+        namespace: str = None,
     ) -> Dict:
         resp = await self.k8s_custom_object_api.patch_namespaced_custom_object(  # type: ignore[attr-defined]
             name=name,
             group=group,
             version=version,
             plural=plural,
-            namespace=self.namespace,
+            namespace=namespace or self.namespace,
             body=body,
             _content_type="application/merge-patch+json",
         )
@@ -153,17 +188,28 @@ class AsyncK8sManager(BaseK8sManager):
         plural: str,
         body: Dict,
         reraise: bool = False,
+        namespace: str = None,
     ) -> Tuple[Dict, bool]:
         try:
             create = await self.create_custom_object(
-                name=name, group=group, version=version, plural=plural, body=body
+                name=name,
+                group=group,
+                version=version,
+                plural=plural,
+                body=body,
+                namespace=namespace,
             )
             return create, True
 
         except ApiException as e_create:
             try:
                 update = await self.update_custom_object(
-                    name=name, group=group, version=version, plural=plural, body=body
+                    name=name,
+                    group=group,
+                    version=version,
+                    plural=plural,
+                    body=body,
+                    namespace=namespace,
                 )
                 return update, False
             except ApiException as e:
@@ -176,7 +222,13 @@ class AsyncK8sManager(BaseK8sManager):
         return {}, False
 
     async def get_custom_object(
-        self, name: str, group: str, version: str, plural: str, reraise: bool = False
+        self,
+        name: str,
+        group: str,
+        version: str,
+        plural: str,
+        reraise: bool = False,
+        namespace: str = None,
     ) -> Optional[Dict]:
         try:
             return await self.k8s_custom_object_api.get_namespaced_custom_object(
@@ -184,7 +236,7 @@ class AsyncK8sManager(BaseK8sManager):
                 group=group,
                 version=version,
                 plural=plural,
-                namespace=self.namespace,
+                namespace=namespace or self.namespace,
             )
         except ApiException as e:
             if reraise:
@@ -192,7 +244,13 @@ class AsyncK8sManager(BaseK8sManager):
             return None
 
     async def delete_custom_object(
-        self, name: str, group: str, version: str, plural: str, reraise: bool = False
+        self,
+        name: str,
+        group: str,
+        version: str,
+        plural: str,
+        reraise: bool = False,
+        namespace: str = None,
     ):
         try:
             await self.k8s_custom_object_api.delete_namespaced_custom_object(
@@ -200,7 +258,7 @@ class AsyncK8sManager(BaseK8sManager):
                 group=group,
                 version=version,
                 plural=plural,
-                namespace=self.namespace,
+                namespace=namespace or self.namespace,
                 body=client.V1DeleteOptions(),
             )
             logger.debug("Custom object `{}` deleted".format(name))
@@ -209,3 +267,18 @@ class AsyncK8sManager(BaseK8sManager):
                 raise PolyaxonK8sError("Connection error: %s" % e) from e
             else:
                 logger.debug("Custom object `{}` was not found".format(name))
+
+    async def list_namespaced_events(
+        self,
+        field_selector: str = None,
+        namespace: str = None,
+        reraise: bool = False,
+        **kwargs
+    ) -> List[client.CoreV1EventList]:
+        return await self._list_namespace_resource(
+            resource_api=self.k8s_api.list_namespaced_event,  # type: ignore[attr-defined]
+            reraise=reraise,
+            namespace=namespace,
+            field_selector=field_selector,
+            **kwargs,
+        )

polyaxon/_k8s/manager/base.py

@@ -28,6 +28,10 @@ class BaseK8sManager:
             instance
         )
 
+    @staticmethod
+    def get_core_polyaxon() -> str:
+        return "app.kubernetes.io/part-of=polyaxon-core"
+
     @classmethod
     def get_config_auth(cls, k8s_config: Optional[Configuration] = None) -> str:
         if not k8s_config or not k8s_config.api_key: