polyaxon 2.1.8__py3-none-any.whl → 2.2.0.post1__py3-none-any.whl

polyaxon/_client/mixin.py

@@ -0,0 +1,39 @@
+from polyaxon import settings
+from polyaxon._client.client import PolyaxonClient
+from polyaxon._schemas.client import ClientConfig
+from polyaxon._utils.fqn_utils import split_owner_team_space
+
+
+class ClientMixin:
+    @property
+    def client(self):
+        if self._client:
+            return self._client
+        self._client = PolyaxonClient()
+        return self._client
+
+    def reset_client(self, **kwargs):
+        if not settings.CLIENT_CONFIG.in_cluster:
+            self._client = PolyaxonClient(
+                ClientConfig.patch_from(settings.CLIENT_CONFIG, **kwargs)
+            )
+
+    @property
+    def owner(self) -> str:
+        return self._owner
+
+    @property
+    def team(self):
+        return self._team
+
+    @property
+    def project(self) -> str:
+        return self._project
+
+    def set_project(self, project: str):
+        self._project = project
+
+    def set_owner(self, owner: str):
+        owner, team = split_owner_team_space(owner)
+        self._owner = owner
+        self._team = team

polyaxon/_client/project.py

@@ -12,6 +12,7 @@ from clipped.utils.validation import validate_tags
 
 from polyaxon._client.client import PolyaxonClient
 from polyaxon._client.decorators import client_handler, get_global_or_inline_config
+from polyaxon._client.mixin import ClientMixin
 from polyaxon._constants.globals import DEFAULT
 from polyaxon._contexts import paths as ctx_paths
 from polyaxon._env_vars.getters.user import get_local_owner
@@ -21,13 +22,17 @@ from polyaxon._sdk.schemas.v1_list_project_versions_response import (
 )
 from polyaxon._sdk.schemas.v1_project import V1Project
 from polyaxon._sdk.schemas.v1_project_version import V1ProjectVersion
-from polyaxon._utils.fqn_utils import get_entity_full_name, get_entity_info
+from polyaxon._utils.fqn_utils import (
+    get_entity_full_name,
+    get_entity_info,
+    split_owner_team_space,
+)
 from polyaxon.exceptions import ApiException, PolyaxonClientException
 from polyaxon.logger import logger
 from traceml.artifacts import V1RunArtifact
 
 
-class ProjectClient:
+class ProjectClient(ClientMixin):
     """ProjectClient is a client to communicate with Polyaxon projects endpoints.
 
     If no values are passed to this class,
@@ -94,26 +99,13 @@ class ProjectClient:
         if not owner:
             raise PolyaxonClientException("Please provide a valid owner.")
 
+        owner, team = split_owner_team_space(owner)
         self._client = client
         self._owner = owner or DEFAULT
+        self._team = team
         self._project = project
         self._project_data = V1Project.construct()
 
-    @property
-    def client(self):
-        if self._client:
-            return self._client
-        self._client = PolyaxonClient()
-        return self._client
-
-    @property
-    def owner(self):
-        return self._owner
-
-    @property
-    def project(self):
-        return self._project
-
     @property
     def project_data(self):
         return self._project_data
@@ -139,11 +131,19 @@ class ProjectClient:
         Returns:
             V1Project, project instance from the response.
         """
-        self._project_data = self.client.projects_v1.create_project(
-            self.owner,
-            data,
-            async_req=False,
-        )
+        if self.team:
+            self._project_data = self.client.projects_v1.create_team_project(
+                self.owner,
+                self.team,
+                data,
+                async_req=False,
+            )
+        else:
+            self._project_data = self.client.projects_v1.create_project(
+                self.owner,
+                data,
+                async_req=False,
+            )
         self._project_data.owner = self.owner
         self._project = self._project_data.name
         return self._project_data

polyaxon/_client/run.py

@@ -35,6 +35,7 @@ from polyaxon import settings
 from polyaxon._cli.errors import handle_cli_error
 from polyaxon._client.client import PolyaxonClient
 from polyaxon._client.decorators import client_handler, get_global_or_inline_config
+from polyaxon._client.mixin import ClientMixin
 from polyaxon._client.store import PolyaxonStore
 from polyaxon._constants.metadata import META_COPY_ARTIFACTS, META_RECOMPILE
 from polyaxon._containers.names import MAIN_CONTAINER_NAMES
@@ -63,7 +64,11 @@ from polyaxon._sdk.schemas.v1_operation_body import V1OperationBody
 from polyaxon._sdk.schemas.v1_project_version import V1ProjectVersion
 from polyaxon._sdk.schemas.v1_run import V1Run
 from polyaxon._sdk.schemas.v1_run_settings import V1RunSettings
-from polyaxon._utils.fqn_utils import get_entity_full_name, to_fqn_name
+from polyaxon._utils.fqn_utils import (
+    get_entity_full_name,
+    split_owner_team_space,
+    to_fqn_name,
+)
 from polyaxon._utils.urls_utils import get_proxy_run_url
 from polyaxon.api import K8S_V1_LOCATION, STREAMS_V1_LOCATION
 from polyaxon.exceptions import ApiException, PolyaxonClientException
@@ -80,7 +85,7 @@ if TYPE_CHECKING:
     from traceml.tracking.run import Run
 
 
-class RunClient:
+class RunClient(ClientMixin):
     """RunClient is a client to communicate with Polyaxon runs endpoints.
 
     If no values are passed to this class,
@@ -146,7 +151,7 @@ class RunClient:
             return
 
         try:
-            owner, project = get_project_or_local(
+            owner, _, project = get_project_or_local(
                 get_entity_full_name(owner=owner, entity=project)
             )
         except PolyaxonClientException:
@@ -166,8 +171,10 @@ class RunClient:
         if error_message and not self._is_offline:
             raise PolyaxonClientException(error_message)
 
+        owner, team = split_owner_team_space(owner)
         self._client = client
         self._owner = owner
+        self._team = team
         self._project = project
         self._run_uuid = (
             get_run_or_local(run_uuid)
@@ -218,12 +225,11 @@ class RunClient:
             return client.config.no_op
         return settings.CLIENT_CONFIG.no_op
 
-    @property
-    def client(self):
-        if self._client:
-            return self._client
-        self._client = PolyaxonClient()
-        return self._client
+    def _use_agent_host(self):
+        if self.settings.agent and self.settings.agent.url:
+            self.reset_client(
+                host=self.settings.agent.url, POLYAXON_HOST=self.settings.agent.url
+            )
 
     @property
     def store(self):
@@ -256,20 +262,6 @@ class RunClient:
             return self.settings.artifacts_store.name
         return None
 
-    @property
-    def owner(self) -> str:
-        return self._owner or ""
-
-    def set_owner(self, owner: str):
-        self._owner = owner
-
-    @property
-    def project(self) -> str:
-        return self._project or ""
-
-    def set_project(self, project: str):
-        self._project = project
-
     @property
     def run_uuid(self) -> str:
         return self._run_uuid
@@ -898,6 +890,7 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
         params = get_logs_params(
             last_file=last_file, last_time=last_time, connection=self.artifacts_store
         )
@@ -921,6 +914,7 @@ class RunClient:
     def inspect(self):
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
         params = get_streams_params(connection=self.artifacts_store, status=self.status)
         return self.client.runs_v1.inspect_run(
             self.namespace, self.owner, self.project, self.run_uuid, **params
@@ -994,6 +988,10 @@ class RunClient:
                 if not container:
                     container = pod_containers[0]
 
+        if not self.settings:
+            self.refresh_data()
+        self._use_agent_host()
+
         url = get_proxy_run_url(
             service=K8S_V1_LOCATION,
             namespace=self.namespace,
@@ -1039,6 +1037,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(self.artifacts_store)
         return self.client.runs_v1.get_run_events(
             self.namespace,
@@ -1074,6 +1074,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(self.artifacts_store)
         return self.client.runs_v1.get_multi_run_events(
             self.namespace,
@@ -1310,6 +1312,7 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
         params = get_streams_params(self.artifacts_store)
         return self.client.runs_v1.get_run_artifact(
             namespace=self.namespace,
@@ -1343,6 +1346,10 @@ class RunClient:
         if not self.run_uuid:
             return
 
+        if not self.settings:
+            self.refresh_data()
+        self._use_agent_host()
+
         lineage_path = lineage.path or ""
         summary = lineage.summary or {}
         is_event = summary.get("is_event")
@@ -1421,6 +1428,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         url = get_proxy_run_url(
             service=STREAMS_V1_LOCATION,
             namespace=self.namespace,
@@ -1460,6 +1469,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         url = get_proxy_run_url(
             service=STREAMS_V1_LOCATION,
             namespace=self.namespace,
@@ -1507,6 +1518,7 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
 
         params = get_streams_params(connection=self.artifacts_store)
         url = get_proxy_run_url(
@@ -1596,6 +1608,7 @@ class RunClient:
 
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
 
         params = get_streams_params(connection=self.artifacts_store)
         url = get_proxy_run_url(
@@ -1626,6 +1639,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(connection=self.artifacts_store)
         self.client.runs_v1.delete_run_artifact(
             namespace=self.namespace,
@@ -1645,6 +1660,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(connection=self.artifacts_store)
         return self.client.runs_v1.delete_run_artifacts(
             namespace=self.namespace,
@@ -1667,6 +1684,8 @@ class RunClient:
         """
         if not self.settings:
             self.refresh_data()
+        self._use_agent_host()
+
         params = get_streams_params(connection=self.artifacts_store)
         return self.client.runs_v1.get_run_artifacts_tree(
             namespace=self.namespace,
@@ -2117,9 +2136,9 @@ class RunClient:
         abspath = filepath if is_abs() else os.path.abspath(filepath)
 
         for_patterns = []
-        if getattr(self, "_artifacts_path"):
+        if getattr(self, "_artifacts_path", None):
             for_patterns.append(getattr(self, "_artifacts_path"))
-        if getattr(self, "_store_path"):
+        if getattr(self, "_store_path", None):
             for_patterns.append(getattr(self, "_store_path"))
         context_root = (
             ctx_paths.CONTEXT_OFFLINE_ROOT

polyaxon/_compiler/contexts/ray_job.py

@@ -33,7 +33,9 @@ class RayJobContextsManager(BaseContextsManager):
                 connection_by_names=connection_by_names,
             )
 
-        return {
+        data = {
             "head": _get_replica(job.head),
-            "workers": {wn: _get_replica(job.workers[wn]) for wn in job.workers},
         }
+        if job.workers:
+            data["workers"] = {wn: _get_replica(job.workers[wn]) for wn in job.workers}
+        return data

polyaxon/_deploy/schemas/proxy.py

@@ -13,4 +13,5 @@ class ProxyConfig(BaseSchemaModel):
     no_proxy: Optional[StrictStr] = Field(alias="noProxy")
     port: Optional[StrictInt]
     host: Optional[StrictStr]
+    protocol: Optional[StrictStr]
     kind: Optional[StrictStr]

polyaxon/_env_vars/getters/owner_entity.py

@@ -5,7 +5,7 @@ from clipped.utils.strings import validate_slug
 
 from polyaxon._constants.globals import DEFAULT
 from polyaxon._env_vars.getters.user import get_local_owner
-from polyaxon._utils.fqn_utils import get_entity_info
+from polyaxon._utils.fqn_utils import get_entity_info, split_owner_team_space
 from polyaxon.exceptions import PolyaxonClientException, PolyaxonSchemaError
 
 
@@ -31,6 +31,8 @@ def resolve_entity_info(entity: str, entity_name: str, is_cli: bool = False):
     if not owner:
         owner = settings.AUTH_CONFIG.username if settings.AUTH_CONFIG else None
 
+    owner, team = split_owner_team_space(owner)
+
     if not all([owner, entity_value]):
         message = "Please provide a valid {}.".format(entity_name)
         if is_cli:
@@ -49,4 +51,4 @@ def resolve_entity_info(entity: str, entity_name: str, is_cli: bool = False):
                 entity_name, entity_value
             )
         )
-    return owner, entity_value
+    return owner, team, entity_value

polyaxon/_env_vars/getters/project.py

@@ -7,7 +7,7 @@ from polyaxon._constants.globals import DEFAULT
 from polyaxon._env_vars.getters.user import get_local_owner
 from polyaxon._managers.project import ProjectConfigManager
 from polyaxon._utils.cache import get_local_project
-from polyaxon._utils.fqn_utils import get_entity_info
+from polyaxon._utils.fqn_utils import get_entity_info, split_owner_team_space
 from polyaxon.exceptions import PolyaxonClientException, PolyaxonSchemaError
 
 
@@ -52,6 +52,8 @@ def get_project_or_local(project=None, is_cli: bool = False):
     if not owner and (not settings.CLI_CONFIG or settings.CLI_CONFIG.is_community):
         owner = DEFAULT
 
+    owner, team = split_owner_team_space(owner)
+
     if not all([owner, project_name]):
         error_message = get_project_error_message(owner, project_name)
         if is_cli:
@@ -79,4 +81,4 @@ def get_project_or_local(project=None, is_cli: bool = False):
             sys.exit(1)
         else:
             raise PolyaxonSchemaError(error_message)
-    return owner, project_name
+    return owner, team, project_name

polyaxon/_env_vars/getters/run.py

@@ -37,9 +37,9 @@ def get_run_or_local(run_uuid=None, is_cli: bool = False):
 
 
 def get_project_run_or_local(project=None, run_uuid=None, is_cli: bool = True):
-    user, project_name = get_project_or_local(project, is_cli=is_cli)
+    owner, team, project_name = get_project_or_local(project, is_cli=is_cli)
     run_uuid = get_run_or_local(run_uuid, is_cli=is_cli)
-    return user, project_name, run_uuid
+    return owner, team, project_name, run_uuid
 
 
 def get_collect_artifacts(arg: Optional[bool] = None, default: Optional[bool] = None):

polyaxon/_env_vars/keys.py

@@ -141,6 +141,7 @@ ENV_KEYS_PROXY_AUTH_USE_RESOLVER = "POLYAXON_PROXY_AUTH_USE_RESOLVER"
 ENV_KEYS_PROXY_HAS_FORWARD_PROXY = "POLYAXON_PROXY_HAS_FORWARD_PROXY"
 ENV_KEYS_PROXY_FORWARD_PROXY_PORT = "POLYAXON_PROXY_FORWARD_PROXY_PORT"
 ENV_KEYS_PROXY_FORWARD_PROXY_HOST = "POLYAXON_PROXY_FORWARD_PROXY_HOST"
+ENV_KEYS_PROXY_FORWARD_PROXY_PROTOCOL = "POLYAXON_PROXY_FORWARD_PROXY_PROTOCOL"
 ENV_KEYS_PROXY_FORWARD_PROXY_KIND = "POLYAXON_PROXY_FORWARD_PROXY_KIND"
 ENV_KEYS_UI_IN_SANDBOX = "POLYAXON_UI_IN_SANDBOX"
 ENV_KEYS_UI_ADMIN_ENABLED = "POLYAXON_UI_ADMIN_ENABLED"

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/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/_local_process/agent.py

@@ -0,0 +1,6 @@
+from polyaxon._local_process.executor import Executor
+from polyaxon._runner.agent.sync_agent import BaseSyncAgent
+
+
+class Agent(BaseSyncAgent):
+    EXECUTOR = Executor

polyaxon/_local_process/converter/__init__.py

@@ -0,0 +1 @@
+

polyaxon/_local_process/converter/base/__init__.py

@@ -0,0 +1 @@
+from polyaxon._local_process.converter.base.base import BaseConverter

polyaxon/_local_process/converter/base/base.py

@@ -0,0 +1,140 @@
+from typing import Dict, Iterable, List, Optional, Union
+
+from polyaxon import settings
+from polyaxon._auxiliaries import V1PolyaxonSidecarContainer
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1Environment, V1Init, V1Plugins
+from polyaxon._k8s import k8s_schemas
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.base.containers import ContainerMixin
+from polyaxon._local_process.converter.base.env_vars import EnvMixin
+from polyaxon._local_process.converter.base.init import InitConverter
+from polyaxon._local_process.converter.base.main import MainConverter
+from polyaxon._local_process.converter.base.mounts import MountsMixin
+from polyaxon._runner.converter import BaseConverter as _BaseConverter
+from polyaxon._runner.kinds import RunnerKind
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+class BaseConverter(
+    MainConverter, InitConverter, ContainerMixin, EnvMixin, MountsMixin, _BaseConverter
+):
+    RUNNER_KIND = RunnerKind.PROCESS
+
+    @classmethod
+    def _get_sidecar_container(
+        cls,
+        container_id: str,
+        polyaxon_sidecar: V1PolyaxonSidecarContainer,
+        env: List[process_types.V1EnvVar],
+        artifacts_store: V1Connection,
+        plugins: V1Plugins,
+        run_path: Optional[str],
+    ) -> Optional[process_types.V1Container]:
+        return None
+
+    @classmethod
+    def _k8s_to_process_env_var(
+        cls,
+        env_var: List[k8s_schemas.V1EnvVar],
+    ) -> List[process_types.V1EnvVar]:
+        if not env_var:
+            return []
+
+        process_env_var = []
+        for item in env_var:
+            if isinstance(item, dict):
+                try:
+                    item = k8s_schemas.V1EnvVar(**item)
+                except (ValueError, TypeError) as e:
+                    raise PolyaxonConverterError(
+                        f"Could not parse env var value `{item}` from the K8S schema in container section"
+                    ) from e
+            process_env_var.append(cls._get_env_var(name=item.name, value=item.value))
+
+        return process_env_var
+
+    @staticmethod
+    def _new_container(name: str) -> process_types.V1Container:
+        return process_types.V1Container(name=name)
+
+    @classmethod
+    def _ensure_container(
+        cls,
+        container: Union[k8s_schemas.V1Container, process_types.V1Container],
+        volumes: List[k8s_schemas.V1Volume],
+    ) -> process_types.V1Container:
+        if not isinstance(container, k8s_schemas.V1Container):
+            return container
+        return process_types.V1Container(
+            name=container.name,
+            command=container.command,
+            args=container.args,
+            env=cls._k8s_to_process_env_var(container.env),
+            working_dir=container.working_dir,
+        )
+
+    def get_replica_resource(
+        self,
+        environment: V1Environment,
+        plugins: V1Plugins,
+        volumes: List[k8s_schemas.V1Volume],
+        init: List[V1Init],
+        sidecars: List[k8s_schemas.V1Container],
+        container: k8s_schemas.V1Container,
+        artifacts_store: V1Connection,
+        connections: List[str],
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        kv_env_vars: List[List],
+        default_sa: Optional[str] = None,
+        ports: List[int] = None,
+    ) -> List[process_types.V1Container]:
+        volumes = volumes or []
+        init = init or []
+        sidecars = sidecars or []
+        connections = connections or []
+        environment = environment or V1Environment()
+        environment.service_account_name = (
+            environment.service_account_name
+            or default_sa
+            or settings.AGENT_CONFIG.runs_sa
+        )
+
+        init_connections = self.filter_connections_from_init(init=init)
+
+        init_containers = self.get_init_containers(
+            polyaxon_init=self.polyaxon_init,
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            init_connections=init_connections,
+            init_containers=self.filter_containers_from_init(init=init),
+            connection_by_names=connection_by_names,
+            log_level=plugins.log_level,
+            volumes=volumes,
+        )
+
+        sidecar_containers = self.get_sidecar_containers(
+            polyaxon_sidecar=self.polyaxon_sidecar,
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            sidecar_containers=sidecars,
+            log_level=plugins.log_level,
+            volumes=volumes,
+        )
+
+        main_container = self.get_main_container(
+            main_container=self._ensure_container(container, volumes=volumes),
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            connections=connections,
+            init_connections=init_connections,
+            connection_by_names=connection_by_names,
+            secrets=secrets,
+            config_maps=config_maps,
+            ports=ports,
+            kv_env_vars=kv_env_vars,
+        )
+
+        return init_containers + sidecar_containers + [main_container]