polyaxon 2.1.0rc9__py3-none-any.whl → 2.6.0__py3-none-any.whl

polyaxon/_flow/run/kubeflow/scheduling_policy.py

@@ -19,8 +19,10 @@ class V1SchedulingPolicy(BaseSchemaModel):
 
     _IDENTIFIER = "schedulingPolicy"
 
-    min_available: Optional[IntOrRef] = Field(alias="minAvailable")
-    queue: Optional[StrictStr]
-    min_resources: Optional[IntOrRef] = Field(alias="minResources")
-    priority_class: Optional[StrictStr] = Field(alias="priorityClass")
-    schedule_timeout_seconds: Optional[IntOrRef] = Field(alias="scheduleTimeoutSeconds")
+    min_available: Optional[IntOrRef] = Field(alias="minAvailable", default=None)
+    queue: Optional[StrictStr] = Field(default=None)
+    min_resources: Optional[IntOrRef] = Field(alias="minResources", default=None)
+    priority_class: Optional[StrictStr] = Field(alias="priorityClass", default=None)
+    schedule_timeout_seconds: Optional[IntOrRef] = Field(
+        alias="scheduleTimeoutSeconds", default=None
+    )

polyaxon/_flow/run/kubeflow/tf_job.py

@@ -169,16 +169,23 @@ class V1TFJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.TFJOB
+    _CUSTOM_DUMP_FIELDS = {"chief", "ps", "worker", "evaluator"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(alias="cleanPodPolicy")
-    scheduling_policy: Optional[V1SchedulingPolicy] = Field(alias="schedulingPolicy")
-    enable_dynamic_worker: Optional[bool] = Field(alias="enableDynamicWorker")
-    success_policy: Optional[str] = Field(alias="successPolicy")
-    chief: Optional[Union[V1KFReplica, RefField]]
-    ps: Optional[Union[V1KFReplica, RefField]]
-    worker: Optional[Union[V1KFReplica, RefField]]
-    evaluator: Optional[Union[V1KFReplica, RefField]]
+    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(
+        alias="cleanPodPolicy", default=None
+    )
+    scheduling_policy: Optional[V1SchedulingPolicy] = Field(
+        alias="schedulingPolicy", default=None
+    )
+    enable_dynamic_worker: Optional[bool] = Field(
+        alias="enableDynamicWorker", default=None
+    )
+    success_policy: Optional[str] = Field(alias="successPolicy", default=None)
+    chief: Optional[Union[V1KFReplica, RefField]] = None
+    ps: Optional[Union[V1KFReplica, RefField]] = None
+    worker: Optional[Union[V1KFReplica, RefField]] = None
+    evaluator: Optional[Union[V1KFReplica, RefField]] = None
 
     def apply_image_destination(self, image: str):
         if self.chief:

polyaxon/_flow/run/kubeflow/xgboost_job.py

@@ -117,12 +117,17 @@ class V1XGBoostJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.XGBJOB
+    _CUSTOM_DUMP_FIELDS = {"master", "worker"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(alias="cleanPodPolicy")
-    scheduling_policy: Optional[V1SchedulingPolicy] = Field(alias="schedulingPolicy")
-    master: Optional[Union[V1KFReplica, RefField]]
-    worker: Optional[Union[V1KFReplica, RefField]]
+    clean_pod_policy: Optional[V1CleanPodPolicy] = Field(
+        alias="cleanPodPolicy", default=None
+    )
+    scheduling_policy: Optional[V1SchedulingPolicy] = Field(
+        alias="schedulingPolicy", default=None
+    )
+    master: Optional[Union[V1KFReplica, RefField]] = None
+    worker: Optional[Union[V1KFReplica, RefField]] = None
 
     def apply_image_destination(self, image: str):
         if self.chief:

polyaxon/_flow/run/ray/ray.py

@@ -135,14 +135,17 @@ class V1RayJob(BaseRun, DestinationImageMixin):
     """
 
     _IDENTIFIER = V1RunKind.RAYJOB
+    _CUSTOM_DUMP_FIELDS = {"head", "workers"}
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    entrypoint: Optional[str]
-    runtime_env: Optional[Union[Dict[str, Any], RefField]] = Field(alias="runtimeEnv")
-    metadata: Optional[Union[Dict[str, str], RefField]]
-    ray_version: Optional[str] = Field(alias="rayVersion")
-    head: Optional[Union[V1RayReplica, RefField]]
-    workers: Optional[Dict[str, Union[V1RayReplica, RefField]]]
+    entrypoint: Optional[str] = None
+    runtime_env: Optional[Union[Dict[str, Any], RefField]] = Field(
+        alias="runtimeEnv", default=None
+    )
+    metadata: Optional[Union[Dict[str, str], RefField]] = Field(default=None)
+    ray_version: Optional[str] = Field(alias="rayVersion", default=None)
+    head: Optional[Union[V1RayReplica, RefField]] = None
+    workers: Optional[Dict[str, Union[V1RayReplica, RefField]]] = Field(default=None)
 
     def apply_image_destination(self, image: str):
         if self.head:

polyaxon/_flow/run/ray/replica.py

@@ -1,6 +1,12 @@
 from typing import Dict, List, Optional, Union
 
-from clipped.compact.pydantic import Field, StrictStr, validator
+from clipped.compact.pydantic import (
+    Field,
+    StrictStr,
+    field_validator,
+    validation_always,
+    validation_before,
+)
 from clipped.types.ref_or_obj import IntOrRef, RefField
 
 from polyaxon._flow.environment import V1Environment
@@ -210,31 +216,34 @@ class V1RayReplica(BaseSchemaModel):
 
     _IDENTIFIER = "replica"
     _SWAGGER_FIELDS = ["volumes", "sidecars", "container"]
-
-    replicas: Optional[IntOrRef]
-    min_replicas: Optional[IntOrRef] = Field(alias="minReplicas")
-    max_replicas: Optional[IntOrRef] = Field(alias="maxReplicas")
-    ray_start_params: Optional[Dict[str, str]] = Field(alias="rayStartParams")
-    environment: Optional[Union[V1Environment, RefField]]
-    connections: Optional[Union[List[StrictStr], RefField]]
-    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]]
-    init: Optional[Union[List[V1Init], RefField]]
-    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]]
-    container: Optional[Union[k8s_schemas.V1Container, RefField]]
-
-    @validator("volumes", always=True, pre=True)
+    _CUSTOM_DUMP_FIELDS = {"environment"}
+
+    replicas: Optional[IntOrRef] = None
+    min_replicas: Optional[IntOrRef] = Field(alias="minReplicas", default=None)
+    max_replicas: Optional[IntOrRef] = Field(alias="maxReplicas", default=None)
+    ray_start_params: Optional[Dict[str, str]] = Field(
+        alias="rayStartParams", default=None
+    )
+    environment: Optional[Union[V1Environment, RefField]] = None
+    connections: Optional[Union[List[StrictStr], RefField]] = None
+    volumes: Optional[Union[List[k8s_schemas.V1Volume], RefField]] = None
+    init: Optional[Union[List[V1Init], RefField]] = None
+    sidecars: Optional[Union[List[k8s_schemas.V1Container], RefField]] = None
+    container: Optional[Union[k8s_schemas.V1Container, RefField]] = None
+
+    @field_validator("volumes", **validation_always, **validation_before)
     def validate_volumes(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_volume(vi) for vi in v]
 
-    @validator("sidecars", always=True, pre=True)
+    @field_validator("sidecars", **validation_always, **validation_before)
     def validate_helper_containers(cls, v):
         if not v:
             return v
         return [k8s_validation.validate_k8s_container(vi) for vi in v]
 
-    @validator("container", always=True, pre=True)
+    @field_validator("container", **validation_always, **validation_before)
     def validate_container(cls, v):
         return k8s_validation.validate_k8s_container(v)
 

polyaxon/_flow/run/resources.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import ClassVar, Optional, Set, Union
 
 from clipped.compact.pydantic import StrictStr
 from clipped.types.numbers import StrictIntOrFloat
@@ -9,18 +9,19 @@ from polyaxon._schemas.base import BaseSchemaModel
 
 
 class V1RunResources(BaseSchemaModel):
-    cpu: Optional[Union[StrictIntOrFloat, StrictStr]]
-    memory: Optional[Union[StrictIntOrFloat, StrictStr]]
-    gpu: Optional[Union[StrictIntOrFloat, StrictStr]]
-    custom: Optional[Union[StrictIntOrFloat, StrictStr]]
-    cost: Optional[Union[StrictIntOrFloat, StrictStr]]
-
-    _MEMORY = "memory"
-    _CPU = "cpu"
-    _GPU = "gpu"
-    _CUSTOM_RESOURCE = "custom"
-    _COST = "cost"
-    _VALUES = {_MEMORY, _CPU, _GPU, _CUSTOM_RESOURCE, _COST}
+    cpu: Optional[Union[StrictIntOrFloat, StrictStr]] = None
+    memory: Optional[Union[StrictIntOrFloat, StrictStr]] = None
+    gpu: Optional[Union[StrictIntOrFloat, StrictStr]] = None
+    custom: Optional[Union[StrictIntOrFloat, StrictStr]] = None
+    cost: Optional[Union[StrictIntOrFloat, StrictStr]] = None
+
+    _MEMORY: ClassVar[str] = "memory"
+    _CPU: ClassVar[str] = "cpu"
+    _GPU: ClassVar[str] = "gpu"
+    _CUSTOM_RESOURCE: ClassVar[str] = "custom"
+    _COST: ClassVar[str] = "cost"
+
+    _VALUES: ClassVar[Set[str]] = {_MEMORY, _CPU, _GPU, _CUSTOM_RESOURCE, _COST}
 
     @classmethod
     def validate_memory(cls, value):

polyaxon/_flow/run/service.py

@@ -272,7 +272,7 @@ class V1Service(V1Job):
     _IDENTIFIER = V1RunKind.SERVICE
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    ports: Optional[Union[List[StrictInt], RefField]]
-    rewrite_path: Optional[BoolOrRef] = Field(alias="rewritePath")
-    is_external: Optional[BoolOrRef] = Field(alias="isExternal")
-    replicas: Optional[IntOrRef]
+    ports: Optional[Union[List[StrictInt], RefField]] = None
+    rewrite_path: Optional[BoolOrRef] = Field(alias="rewritePath", default=None)
+    is_external: Optional[BoolOrRef] = Field(alias="isExternal", default=None)
+    replicas: Optional[IntOrRef] = None

polyaxon/_flow/schedules/cron.py

@@ -114,7 +114,7 @@ class V1CronSchedule(BaseSchemaModel):
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
     cron: StrictStr
-    start_at: Optional[DatetimeOrRef] = Field(alias="startAt")
-    end_at: Optional[DatetimeOrRef] = Field(alias="endAt")
-    max_runs: Optional[IntOrRef] = Field(alias="maxRuns")
-    depends_on_past: Optional[BoolOrRef] = Field(alias="dependsOnPast")
+    start_at: Optional[DatetimeOrRef] = Field(alias="startAt", default=None)
+    end_at: Optional[DatetimeOrRef] = Field(alias="endAt", default=None)
+    max_runs: Optional[IntOrRef] = Field(alias="maxRuns", default=None)
+    depends_on_past: Optional[BoolOrRef] = Field(alias="dependsOnPast", default=None)

polyaxon/_flow/schedules/interval.py

@@ -117,8 +117,8 @@ class V1IntervalSchedule(BaseSchemaModel):
     _USE_DISCRIMINATOR = True
 
     kind: Literal[_IDENTIFIER] = _IDENTIFIER
-    start_at: Optional[DatetimeOrRef] = Field(alias="startAt")
-    end_at: Optional[DatetimeOrRef] = Field(alias="endAt")
-    max_runs: Optional[IntOrRef] = Field(alias="maxRuns")
+    start_at: Optional[DatetimeOrRef] = Field(alias="startAt", default=None)
+    end_at: Optional[DatetimeOrRef] = Field(alias="endAt", default=None)
+    max_runs: Optional[IntOrRef] = Field(alias="maxRuns", default=None)
     frequency: TimeDeltaOrRef
-    depends_on_past: Optional[BoolOrRef] = Field(alias="dependsOnPast")
+    depends_on_past: Optional[BoolOrRef] = Field(alias="dependsOnPast", default=None)

polyaxon/_flow/templates/__init__.py

@@ -80,9 +80,9 @@ class V1Template(BaseSchemaModel):
 
     _IDENTIFIER = "template"
 
-    enabled: Optional[BoolOrRef]
-    description: Optional[StrictStr]
-    fields: Optional[Union[List[StrictStr], RefField]]
+    enabled: Optional[BoolOrRef] = None
+    description: Optional[StrictStr] = None
+    fields: Optional[Union[List[StrictStr], RefField]] = None
 
 
 class TemplateMixinConfig:

polyaxon/_flow/termination/__init__.py

@@ -99,6 +99,6 @@ class V1Termination(BaseSchemaModel):
 
     _IDENTIFIER = "termination"
 
-    max_retries: Optional[IntOrRef] = Field(alias="maxRetries")
-    ttl: Optional[IntOrRef]
-    timeout: Optional[IntOrRef]
+    max_retries: Optional[IntOrRef] = Field(alias="maxRetries", default=None)
+    ttl: Optional[IntOrRef] = None
+    timeout: Optional[IntOrRef] = None

polyaxon/_fs/async_manager.py

@@ -314,7 +314,7 @@ async def list_files(
 
 
 async def delete_file_or_dir(
-    fs: FSSystem, store_path: str, subpath: Union[str], is_file: bool
+    fs: FSSystem, store_path: str, subpath: Optional[str], is_file: bool
 ) -> bool:
     try:
         await ensure_async_execution(

polyaxon/_fs/fs.py

@@ -56,7 +56,7 @@ def _get_fs_from_connection(
 
 async def get_async_fs_from_connection(connection: Optional[V1Connection], **kwargs):
     fs = _get_fs_from_connection(connection=connection, asynchronous=True, **kwargs)
-    if fs.async_impl and hasattr(fs, "set_session"):
+    if hasattr(fs, "async_impl") and hasattr(fs, "set_session"):
         await fs.set_session()
     return fs
 

polyaxon/_fs/watcher.py

@@ -3,36 +3,43 @@ import os
 from datetime import datetime
 from typing import Dict, List, Optional, Set, Tuple
 
+from clipped.compact.pydantic import PYDANTIC_VERSION
 from clipped.utils.dates import path_last_modified
 from clipped.utils.paths import get_files_and_dirs_in_path
 
 from polyaxon._contexts import paths as ctx_paths
-from polyaxon._schemas.base import BaseSchemaModel
+from polyaxon._schemas.base import BaseSchemaModel, RootModel
 
 
-class PathData(BaseSchemaModel):
-    __root__: Tuple[str, datetime, str]
+class PathData(RootModel):
+    if PYDANTIC_VERSION.startswith("2."):
+        root: Tuple[str, datetime, str]
+    else:
+        __root__: Tuple[str, datetime, str]
 
-    class Config(BaseSchemaModel.Config):
+    class Config:
         validate_assignment = False
 
     @property
     def base(self) -> str:
-        return self.__root__[0]
+        return self.get_root()[0]
 
     @property
     def ts(self) -> datetime:
-        if isinstance(self.__root__[1], str):
-            self.__root__ = (
-                self.__root__[0],
-                datetime.fromisoformat(self.__root__[1]),
-                self.__root__[1],
+        root = self.get_root()
+        if isinstance(root[1], str):
+            self.set_root(
+                (
+                    root[0],
+                    datetime.fromisoformat(root[1]),
+                    root[1],
+                )
             )
-        return self.__root__[1]
+        return root[1]
 
     @property
     def op(self) -> str:
-        return self.__root__[2]
+        return self.get_root()[2]
 
 
 class FSWatcher(BaseSchemaModel):
@@ -42,8 +49,8 @@ class FSWatcher(BaseSchemaModel):
     _RM = "rm"
     _NOOP = ""
 
-    dir_mapping: Optional[Dict[str, PathData]]
-    file_mapping: Optional[Dict[str, PathData]]
+    dir_mapping: Optional[Dict[str, PathData]] = None
+    file_mapping: Optional[Dict[str, PathData]] = None
 
     @property
     def dirs_mp(self) -> Dict[str, PathData]:
@@ -71,17 +78,11 @@ class FSWatcher(BaseSchemaModel):
         data = mapping.get(rel_path)
         if data:
             if current_ts > data.ts:
-                mapping[rel_path] = PathData.construct(
-                    __root__=(base_path, current_ts, self._PUT)
-                )
+                mapping[rel_path] = PathData.make((base_path, current_ts, self._PUT))
             else:
-                mapping[rel_path] = PathData.construct(
-                    __root__=(base_path, data.ts, self._NOOP)
-                )
+                mapping[rel_path] = PathData.make((base_path, data.ts, self._NOOP))
         else:
-            mapping[rel_path] = PathData.construct(
-                __root__=(base_path, current_ts, self._PUT)
-            )
+            mapping[rel_path] = PathData.make((base_path, current_ts, self._PUT))
         return mapping
 
     def sync_file(self, path: str, base_path: str):
@@ -92,12 +93,10 @@ class FSWatcher(BaseSchemaModel):
 
     def init(self):
         self.dir_mapping = {
-            p: PathData.construct(__root__=(d.base, d.ts, self._RM))
-            for p, d in self.dirs_mp.items()
+            p: PathData.make((d.base, d.ts, self._RM)) for p, d in self.dirs_mp.items()
         }
         self.file_mapping = {
-            p: PathData.construct(__root__=(d.base, d.ts, self._RM))
-            for p, d in self.files_mp.items()
+            p: PathData.make((d.base, d.ts, self._RM)) for p, d in self.files_mp.items()
         }
 
     def sync(self, path: str, exclude: Optional[List[str]] = None):

polyaxon/_k8s/converter/base/base.py

@@ -2,6 +2,7 @@ import copy
 
 from typing import Dict, Iterable, List, Optional
 
+from clipped.utils.enums import get_enum_value
 from clipped.utils.sanitizers import sanitize_string_dict
 from clipped.utils.strings import slugify
 
@@ -85,7 +86,7 @@ class BaseConverter(
             "operation.polyaxon.com/name": self.run_name,
             "operation.polyaxon.com/owner": self.owner_name,
             "operation.polyaxon.com/project": self.project_name,
-            "operation.polyaxon.com/kind": self.K8S_ANNOTATIONS_KIND,
+            "operation.polyaxon.com/kind": get_enum_value(self.K8S_ANNOTATIONS_KIND),
         }
 
     def get_annotations(

polyaxon/_k8s/converter/base/main.py

@@ -81,6 +81,7 @@ class MainConverter(_BaseConverter):
             secrets=requested_secrets,
             config_maps=requested_config_maps,
         )
+        env += self._get_resources_env_vars(main_container.resources)
 
         # Env from
         env_from = self._get_env_from_k8s_resources(

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/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/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/k8s_validation.py

@@ -14,7 +14,7 @@ def _validate_schema(value: Optional[Union[Swagger, Dict]], cls: Type[Swagger]):
         return cls(**{to_snake_case(k): value[k] for k in value})
     if isinstance(value, cls):
         return value
-    raise TypeError(
+    raise ValueError(
         "This field expects a dict or an instance of {}.".format(cls.__name__)
     )