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

polyaxon/_k8s/converter/base/sidecar.py

@@ -33,15 +33,26 @@ class SidecarConverter(_BaseConverter):
 
     @staticmethod
     def _get_sidecar_args(
-        container_id: str, sleep_interval: int, sync_interval: int, monitor_logs: bool
+        container_id: str,
+        sleep_interval: int,
+        sync_interval: int,
+        monitor_logs: bool,
+        monitor_spec: bool,
     ) -> List[str]:
         args = [
             "--container-id={}".format(container_id),
             "--sleep-interval={}".format(sleep_interval),
             "--sync-interval={}".format(sync_interval),
         ]
+        # enable monitor logs and spec by default
+        if monitor_logs is None:
+            monitor_logs = True
+        if monitor_spec is None:
+            monitor_spec = True
         if monitor_logs:
             args.append("--monitor-logs")
+        if monitor_spec:
+            args.append("--monitor-spec")
         return args
 
     @classmethod
@@ -87,6 +98,7 @@ class SidecarConverter(_BaseConverter):
         sleep_interval = polyaxon_sidecar.sleep_interval
         sync_interval = polyaxon_sidecar.sync_interval
         monitor_logs = polyaxon_sidecar.monitor_logs
+        monitor_spec = polyaxon_sidecar.monitor_spec
         if plugins and plugins.sidecar:
             if plugins.sidecar.sleep_interval:
                 sleep_interval = plugins.sidecar.sleep_interval
@@ -94,11 +106,14 @@ class SidecarConverter(_BaseConverter):
                 sync_interval = plugins.sidecar.sync_interval
             if plugins.sidecar.monitor_logs:
                 monitor_logs = plugins.sidecar.monitor_logs
+            if plugins.sidecar.monitor_spec:
+                monitor_spec = plugins.sidecar.monitor_spec
         sidecar_args = cls._get_sidecar_args(
             container_id=container_id,
             sleep_interval=sleep_interval,
             sync_interval=sync_interval,
             monitor_logs=monitor_logs,
+            monitor_spec=monitor_spec,
         )
 
         env_from = []

polyaxon/_k8s/converter/common/accelerators.py

@@ -37,14 +37,17 @@ def requests_gpu(resources: k8s_schemas.V1ResourceRequirements) -> bool:
     if not resources:
         return False
 
+    if not isinstance(resources, k8s_schemas.V1ResourceRequirements):
+        resources = k8s_schemas.V1ResourceRequirements(**resources)
+
     if resources.requests:
-        for key in resources.requests.keys():
-            if "gpu" in key:
+        for key, val in resources.requests.items():
+            if "gpu" in key and val is not None and val > 0:
                 return True
 
     if resources.limits:
-        for key in resources.limits.keys():
-            if "gpu" in key:
+        for key, val in resources.limits.items():
+            if "gpu" in key and val is not None and val > 0:
                 return True
 
     return False

polyaxon/_k8s/converter/converters/job.py

@@ -39,7 +39,7 @@ class JobConverter(JobMixin, BaseConverter):
             default_sa=default_sa,
         )
         return get_job_custom_resource(
-            namespace=self.namespace,
+            namespace=compiled_operation.namespace or self.namespace,
             main_container=replica_spec.main_container,
             sidecar_containers=replica_spec.sidecar_containers,
             init_containers=replica_spec.init_containers,

polyaxon/_k8s/converter/converters/kubeflow/paddle_job.py

@@ -58,6 +58,7 @@ class PaddleJobConverter(PaddleJobMixin, BaseConverter):
             termination=compiled_operation.termination,
             clean_pod_policy=job.clean_pod_policy,
             scheduling_policy=job.scheduling_policy,
+            elastic_policy=job.elastic_policy,
             collect_logs=plugins.collect_logs,
             sync_statuses=plugins.sync_statuses,
             notifications=plugins.notifications,

polyaxon/_k8s/converter/converters/kubeflow/pytroch_job.py

@@ -58,6 +58,8 @@ class PytorchJobConverter(PytorchJobMixin, BaseConverter):
             termination=compiled_operation.termination,
             clean_pod_policy=job.clean_pod_policy,
             scheduling_policy=job.scheduling_policy,
+            elastic_policy=job.elastic_policy,
+            n_proc_per_node=job.n_proc_per_node,
             collect_logs=plugins.collect_logs,
             sync_statuses=plugins.sync_statuses,
             notifications=plugins.notifications,

polyaxon/_k8s/converter/converters/kubeflow/tf_job.py

@@ -63,6 +63,7 @@ class TfJobConverter(TFJobMixin, BaseConverter):
             collect_logs=plugins.collect_logs,
             clean_pod_policy=job.clean_pod_policy,
             scheduling_policy=job.scheduling_policy,
+            success_policy=job.success_policy,
             enable_dynamic_worker=job.enable_dynamic_worker,
             sync_statuses=plugins.sync_statuses,
             notifications=plugins.notifications,

polyaxon/_k8s/converter/converters/ray_job.py

@@ -58,7 +58,9 @@ class RayJobConverter(RayJobMixin, BaseConverter):
             config=compiled_operation.plugins, auth=default_auth
         )
         head = _get_replica(job.head)
-        workers = {n: _get_replica(w) for n, w in job.workers.items()}
+        workers = None
+        if job.workers:
+            workers = {n: _get_replica(w) for n, w in job.workers.items()}
         labels = self.get_labels(version=pkg.VERSION, labels={})
 
         return get_ray_job_custom_resource(
@@ -68,7 +70,7 @@ class RayJobConverter(RayJobMixin, BaseConverter):
             workers=workers,
             entrypoint=job.entrypoint,
             metadata=job.metadata,
-            runtime_env=encode(orjson_dumps(job.runtime_env)),
+            runtime_env=orjson_dumps(job.runtime_env),
             ray_version=job.ray_version,
             termination=compiled_operation.termination,
             collect_logs=plugins.collect_logs,

polyaxon/_k8s/custom_resources/dask_job.py

@@ -46,6 +46,9 @@ def get_dask_replicas_template(
     if liveness_probe and replica.main_container.liveness_probe is None:
         replica.main_container.liveness_probe = liveness_probe
 
+    labels = {**labels, **replica.labels}
+    annotations = {**annotations, **replica.annotations}
+
     metadata, pod_spec = get_pod_spec(
         namespace=namespace,
         main_container=replica.main_container,

polyaxon/_k8s/custom_resources/kubeflow/common.py

@@ -16,6 +16,9 @@ def get_kf_replicas_template(
     if not replica:
         return
 
+    labels = {**labels, **replica.labels}
+    annotations = {**annotations, **replica.annotations}
+
     metadata, pod_spec = get_pod_spec(
         namespace=namespace,
         main_container=replica.main_container,
@@ -30,6 +33,6 @@ def get_kf_replicas_template(
 
     template_spec[replica_name] = {
         "replicas": replica.num_replicas,
-        "restartPolicy": pod_spec.restart_policy or "Never",
+        "restartPolicy": pod_spec.restart_policy or "OnFailure",
         "template": get_pod_template_spec(metadata=metadata, pod_spec=pod_spec),
     }

polyaxon/_k8s/custom_resources/kubeflow/paddle_job.py

@@ -1,6 +1,11 @@
 from typing import Dict, List, Optional
 
-from polyaxon._flow import V1Notification, V1SchedulingPolicy, V1Termination
+from polyaxon._flow import (
+    V1Notification,
+    V1PaddleElasticPolicy,
+    V1SchedulingPolicy,
+    V1Termination,
+)
 from polyaxon._k8s.custom_resources.kubeflow.common import get_kf_replicas_template
 from polyaxon._k8s.custom_resources.operation import get_operation_custom_object
 from polyaxon._k8s.custom_resources.setter import (
@@ -25,6 +30,7 @@ def get_paddle_job_custom_resource(
     notifications: List[V1Notification],
     clean_pod_policy: Optional[str],
     scheduling_policy: Optional[V1SchedulingPolicy],
+    elastic_policy: Optional[V1PaddleElasticPolicy],
     labels: Dict[str, str],
     annotations: Dict[str, str],
 ) -> Dict:
@@ -58,6 +64,9 @@ def get_paddle_job_custom_resource(
         template_spec=template_spec, scheduling_policy=scheduling_policy
     )
 
+    if elastic_policy:
+        template_spec["elasticPolicy"] = elastic_policy.to_light_dict()
+
     custom_object = {"paddleJobSpec": template_spec}
     custom_object = set_termination(
         custom_object=custom_object, termination=termination

polyaxon/_k8s/custom_resources/kubeflow/pytorch_job.py

@@ -1,6 +1,11 @@
 from typing import Dict, List, Optional
 
-from polyaxon._flow import V1Notification, V1SchedulingPolicy, V1Termination
+from polyaxon._flow import (
+    V1Notification,
+    V1PytorchElasticPolicy,
+    V1SchedulingPolicy,
+    V1Termination,
+)
 from polyaxon._k8s.custom_resources.kubeflow.common import get_kf_replicas_template
 from polyaxon._k8s.custom_resources.operation import get_operation_custom_object
 from polyaxon._k8s.custom_resources.setter import (
@@ -25,6 +30,8 @@ def get_pytorch_job_custom_resource(
     notifications: List[V1Notification],
     clean_pod_policy: Optional[str],
     scheduling_policy: Optional[V1SchedulingPolicy],
+    elastic_policy: Optional[V1PytorchElasticPolicy],
+    n_proc_per_node: Optional[int],
     labels: Dict[str, str],
     annotations: Dict[str, str],
 ) -> Dict:
@@ -58,6 +65,12 @@ def get_pytorch_job_custom_resource(
         template_spec=template_spec, scheduling_policy=scheduling_policy
     )
 
+    if elastic_policy:
+        template_spec["elasticPolicy"] = elastic_policy.to_light_dict()
+
+    if n_proc_per_node is not None:
+        template_spec["nProcPerNode"] = str(n_proc_per_node)
+
     custom_object = {"pytorchJobSpec": template_spec}
     custom_object = set_termination(
         custom_object=custom_object, termination=termination

polyaxon/_k8s/custom_resources/kubeflow/tf_job.py

@@ -28,6 +28,7 @@ def get_tf_job_custom_resource(
     clean_pod_policy: Optional[str],
     scheduling_policy: Optional[V1SchedulingPolicy],
     enable_dynamic_worker: bool,
+    success_policy: Optional[str],
     labels: Dict[str, str],
     annotations: Dict[str, str],
 ) -> Dict:
@@ -73,6 +74,9 @@ def get_tf_job_custom_resource(
     if enable_dynamic_worker:
         template_spec["enableDynamicWorker"] = enable_dynamic_worker
 
+    if success_policy:
+        template_spec["successPolicy"] = success_policy
+
     template_spec = {"replicaSpecs": template_spec}
 
     template_spec = set_clean_pod_policy(

polyaxon/_k8s/custom_resources/ray_job.py

@@ -26,6 +26,9 @@ def _get_ray_replicas_template(
     if not replica:
         return
 
+    labels = {**labels, **replica.labels}
+    annotations = {**annotations, **replica.annotations}
+
     metadata, pod_spec = get_pod_spec(
         namespace=namespace,
         main_container=replica.main_container,

polyaxon/_k8s/custom_resources/setter.py

@@ -46,7 +46,7 @@ def set_notify(custom_object: Dict, notifications: List[V1Notification]) -> Dict
 def set_clean_pod_policy(template_spec: Dict, clean_pod_policy: str) -> Dict:
     if not clean_pod_policy:
         # Sets default clean pod policy
-        clean_pod_policy = "All"
+        clean_pod_policy = "None"
 
     template_spec["cleanPodPolicy"] = clean_pod_policy.capitalize()
     return template_spec

polyaxon/_k8s/executor/async_executor.py

@@ -14,6 +14,8 @@ class AsyncExecutor(BaseExecutor):
         )
 
     async def refresh(self):
+        if self._manager:
+            await self._manager.close()
         manager = super().refresh()
         await manager.setup()
         return manager

polyaxon/_k8s/executor/base.py

@@ -5,7 +5,7 @@ from kubernetes.client import Configuration
 
 from polyaxon import settings
 from polyaxon._k8s.converter.converters import CONVERTERS
-from polyaxon._k8s.converter.mixins import MIXIN_MAPPING
+from polyaxon._k8s.converter.mixins import MIXIN_MAPPING, BaseMixin
 from polyaxon._runner.executor import BaseExecutor as _BaseExecutor
 from polyaxon._runner.kinds import RunnerKind
 from polyaxon._utils.fqn_utils import get_resource_name
@@ -56,7 +56,9 @@ class BaseExecutor(_BaseExecutor):
         api = k8s_client.ApiClient()
         return api.sanitize_for_serialization(resource)
 
-    def create(self, run_uuid: str, run_kind: str, resource: Dict) -> Dict:
+    def create(
+        self, run_uuid: str, run_kind: str, resource: Dict, namespace: str = None
+    ) -> Dict:
         mixin = self._get_mixin_for_kind(kind=run_kind)
         resource_name = get_resource_name(run_uuid)
         return self.manager.create_custom_object(
@@ -65,9 +67,12 @@ class BaseExecutor(_BaseExecutor):
             version=mixin.API_VERSION,
             plural=mixin.PLURAL,
             body=resource,
+            namespace=namespace,
         )
 
-    def apply(self, run_uuid: str, run_kind: str, resource: Dict) -> Dict:
+    def apply(
+        self, run_uuid: str, run_kind: str, resource: Dict, namespace: str = None
+    ) -> Dict:
         mixin = self._get_mixin_for_kind(kind=run_kind)
         resource_name = get_resource_name(run_uuid)
         return self.manager.update_custom_object(
@@ -76,9 +81,10 @@ class BaseExecutor(_BaseExecutor):
             version=mixin.API_VERSION,
             plural=mixin.PLURAL,
             body=resource,
+            namespace=namespace,
         )
 
-    def stop(self, run_uuid: str, run_kind: str):
+    def stop(self, run_uuid: str, run_kind: str, namespace: str = None):
         mixin = self._get_mixin_for_kind(kind=run_kind)
         resource_name = get_resource_name(run_uuid)
         return self.manager.delete_custom_object(
@@ -86,16 +92,18 @@ class BaseExecutor(_BaseExecutor):
             group=mixin.GROUP,
             version=mixin.API_VERSION,
             plural=mixin.PLURAL,
+            namespace=namespace,
         )
 
-    def clean(self, run_uuid: str, run_kind: str):
+    def clean(self, run_uuid: str, run_kind: str, namespace: str = None):
         return self.apply(
             run_uuid=run_uuid,
             run_kind=run_kind,
             resource={"metadata": {"finalizers": None}},
+            namespace=namespace,
         )
 
-    def get(self, run_uuid: str, run_kind: str):
+    def get(self, run_uuid: str, run_kind: str, namespace: str = None):
         mixin = self._get_mixin_for_kind(kind=run_kind)
         resource_name = get_resource_name(run_uuid)
         return self.manager.get_custom_object(
@@ -103,4 +111,13 @@ class BaseExecutor(_BaseExecutor):
             group=mixin.GROUP,
             version=mixin.API_VERSION,
             plural=mixin.PLURAL,
+            namespace=namespace,
+        )
+
+    def list_ops(self, namespace: str = None):
+        return self.manager.list_custom_objects(
+            group=BaseMixin.GROUP,
+            version=BaseMixin.API_VERSION,
+            plural=BaseMixin.PLURAL,
+            namespace=namespace,
         )

polyaxon/_k8s/logging/async_monitor.py

@@ -6,6 +6,7 @@ from clipped.utils.tz import now
 from kubernetes_asyncio.client.models import V1Pod
 from kubernetes_asyncio.client.rest import ApiException
 
+from polyaxon._flow import V1RunKind
 from polyaxon._k8s.manager.async_manager import AsyncK8sManager
 from traceml.logging import V1Log, V1Logs
 
@@ -65,7 +66,7 @@ async def query_k8s_operation_logs(
     new_time = now()
     params = {}
     if last_time:
-        since_seconds = (new_time - last_time).total_seconds() - 1
+        since_seconds = (new_time - last_time).total_seconds()
         params["since_seconds"] = int(since_seconds)
     if stream:
         params["tail_lines"] = V1Logs._CHUNK_SIZE
@@ -82,9 +83,28 @@ async def query_k8s_operation_logs(
             **params,
         )
 
+    if logs and last_time:
+        # make sure to filter logs larger than last_time
+        logs = [log for log in logs if log.timestamp > last_time]
+    if logs and logs[-1].timestamp:
+        new_time = logs[-1].timestamp
     return logs, new_time
 
 
+async def collect_agent_service_logs(
+    k8s_manager: AsyncK8sManager, pod: V1Pod
+) -> List[V1Log]:
+    if not pod or not pod.spec.containers:
+        return []
+    container = pod.spec.containers[0]
+    return await handle_container_logs(
+        k8s_manager=k8s_manager,
+        pod=pod,
+        container_name=container.name,
+        tail_lines=V1Logs._CHUNK_SIZE,
+    )
+
+
 async def query_k8s_pod_logs(
     k8s_manager: AsyncK8sManager,
     pod: V1Pod,
@@ -94,13 +114,138 @@ async def query_k8s_pod_logs(
     new_time = now()
     params = {}
     if last_time:
-        since_seconds = (new_time - last_time).total_seconds() - 1
+        since_seconds = (new_time - last_time).total_seconds()
         params["since_seconds"] = int(since_seconds)
     if stream:
         params["tail_lines"] = V1Logs._CHUNK_SIZE
 
     logs = await handle_pod_logs(k8s_manager=k8s_manager, pod=pod, **params)
 
-    if logs:
-        last_time = logs[-1].timestamp
-    return logs, last_time
+    if logs and last_time:
+        # make sure to filter logs larger than last_time
+        logs = [log for log in logs if log.timestamp > last_time]
+    if logs and logs[-1].timestamp:
+        new_time = logs[-1].timestamp
+    return logs, new_time
+
+
+async def get_op_pods_and_services(
+    k8s_manager: AsyncK8sManager,
+    run_uuid: str,
+    run_kind: str,
+):
+    pods = await k8s_manager.list_pods(
+        label_selector=k8s_manager.get_managed_by_polyaxon(run_uuid)
+    )
+    services = []
+    if V1RunKind.has_service(run_kind):
+        services = await k8s_manager.list_services(
+            label_selector=k8s_manager.get_managed_by_polyaxon(run_uuid)
+        )
+
+    return pods, services
+
+
+async def get_resource_events(
+    k8s_manager: AsyncK8sManager, resource_type: str, resource_name: str
+):
+    field_selector = (
+        f"involvedObject.kind={resource_type},involvedObject.name={resource_name}"
+    )
+    try:
+        events = await k8s_manager.list_namespaced_events(field_selector=field_selector)
+
+        all_events = []
+        for event in events:
+            event_data = {
+                "reason": event.reason,
+                "message": event.message,
+                "first_timestamp": event.first_timestamp,
+                "last_timestamp": event.last_timestamp,
+                "count": event.count,
+                "type": event.type,
+            }
+            all_events.append(event_data)
+
+        return all_events
+
+    except ApiException as e:
+        print(f"Exception when calling CoreV1Api->list_namespaced_event: {e}")
+        return []
+
+
+async def get_op_spec(
+    k8s_manager: AsyncK8sManager,
+    run_uuid: str,
+    run_kind: str,
+):
+    pods, services = await get_op_pods_and_services(
+        k8s_manager=k8s_manager,
+        run_uuid=run_uuid,
+        run_kind=run_kind,
+    )
+    pods_list = {}
+    for pod in pods or []:
+        pods_list[
+            pod.metadata.name
+        ] = k8s_manager.api_client.sanitize_for_serialization(pod)
+        pods_list[pod.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Pod",
+            resource_name=pod.metadata.name,
+        )
+    services_list = {}
+    for service in services or []:
+        services_list[
+            service.metadata.name
+        ] = k8s_manager.api_client.sanitize_for_serialization(service)
+        services_list[service.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Service",
+            resource_name=service.metadata.name,
+        )
+    data = {"pods": pods_list, "services": services_list}
+    return data, pods, services
+
+
+async def get_agent_pods_and_services(
+    k8s_manager: AsyncK8sManager,
+):
+    pods = await k8s_manager.list_pods(
+        label_selector=k8s_manager.get_core_polyaxon(),
+    )
+    services = await k8s_manager.list_services(
+        label_selector=k8s_manager.get_core_polyaxon(),
+    )
+    return pods, services
+
+
+async def get_agent_spec(
+    k8s_manager: AsyncK8sManager,
+):
+    pods, services = await get_agent_pods_and_services(
+        k8s_manager=k8s_manager,
+    )
+    pods_list = {}
+    for pod in pods or []:
+        pods_list[
+            pod.metadata.name
+        ] = k8s_manager.api_client.sanitize_for_serialization(pod)
+        pods_list[pod.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Pod",
+            resource_name=pod.metadata.name,
+        )
+    data = {"pods": pods_list}
+    services_list = {}
+    for service in services or []:
+        services_list[
+            service.metadata.name
+        ] = k8s_manager.api_client.sanitize_for_serialization(service)
+        services_list[service.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Service",
+            resource_name=service.metadata.name,
+        )
+    data["services"] = services_list
+    return data, pods, services