Skip to content

Commit 8c37856

Browse files
polyaxon-cipolyaxon-ci
committed
Add logic for dask and ray clusters
1 parent 49edb10 commit 8c37856

47 files changed

Lines changed: 336 additions & 347 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎cli/polyaxon/_client/run.py‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -906,7 +906,10 @@ def get_logs(self, last_file=None, last_time=None) -> "V1Logs":
906906
self.refresh_data()
907907
self._use_agent_host()
908908
params = get_logs_params(
909-
last_file=last_file, last_time=last_time, connection=self.artifacts_store
909+
last_file=last_file,
910+
last_time=last_time,
911+
connection=self.artifacts_store,
912+
kind=self.run_data.kind,
910913
)
911914
return self.client.runs_v1.get_run_logs(
912915
self.namespace, self.owner, self.project, self.run_uuid, **params
@@ -931,7 +934,12 @@ def inspect(self):
931934
self._use_agent_host()
932935
params = get_streams_params(connection=self.artifacts_store, status=self.status)
933936
return self.client.runs_v1.inspect_run(
934-
self.namespace, self.owner, self.project, self.run_uuid, **params
937+
self.namespace,
938+
self.owner,
939+
self.project,
940+
self.run_uuid,
941+
self.run_data.kind,
942+
**params,
935943
)
936944

937945
@client_handler(check_no_op=True, check_offline=True)
@@ -2627,6 +2635,7 @@ def list(
26272635
sort: Optional[str] = None,
26282636
limit: Optional[int] = None,
26292637
offset: Optional[int] = None,
2638+
all_projects: bool = False,
26302639
):
26312640
"""Lists runs under the current owner - project.
26322641
@@ -2639,13 +2648,18 @@ def list(
26392648
[Run PQL](/docs/core/query-language/runs/#sort)
26402649
limit: int, optional, limit of runs to return.
26412650
offset: int, optional, offset pages to paginate runs.
2651+
all_projects: bool, optional, default: False
26422652
26432653
Returns:
26442654
List[V1Run], list of run instances.
26452655
"""
26462656
params = get_query_params(
26472657
limit=limit or 20, offset=offset, query=query, sort=sort
26482658
)
2659+
if all_projects:
2660+
return self.client.organizations_v1.get_organization_runs(
2661+
self.owner, **params
2662+
)
26492663
return self.client.runs_v1.list_runs(self.owner, self.project, **params)
26502664

26512665
@client_handler(check_no_op=True, check_offline=True)

‎cli/polyaxon/_compiler/contexts/contexts.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,14 @@
33
from datetime import datetime
44
from typing import Dict, Optional
55

6-
from polyaxon._compiler.contexts.dask_job import DaskJobContextsManager
6+
from polyaxon._compiler.contexts.dask_cluster import DaskClusterContextsManager
77
from polyaxon._compiler.contexts.job import JobContextsManager
88
from polyaxon._compiler.contexts.kubeflow import (
99
MPIJobContextsManager,
1010
PytorchJobContextsManager,
1111
TfJobContextsManager,
1212
)
13-
from polyaxon._compiler.contexts.ray_job import RayJobContextsManager
13+
from polyaxon._compiler.contexts.ray_cluster import RayClusterContextsManager
1414
from polyaxon._compiler.contexts.service import ServiceContextsManager
1515
from polyaxon._connections import V1Connection
1616
from polyaxon._contexts import keys as ctx_keys
@@ -30,8 +30,8 @@
3030
V1RunKind.MPIJOB: MPIJobContextsManager,
3131
V1RunKind.TFJOB: TfJobContextsManager,
3232
V1RunKind.PYTORCHJOB: PytorchJobContextsManager,
33-
V1RunKind.RAYJOB: RayJobContextsManager,
34-
V1RunKind.DASKJOB: DaskJobContextsManager,
33+
V1RunKind.RAYCLUSTER: RayClusterContextsManager,
34+
V1RunKind.DASKCLUSTER: DaskClusterContextsManager,
3535
}
3636

3737

cli/polyaxon/_compiler/contexts/dask_job.py renamed to cli/polyaxon/_compiler/contexts/dask_cluster.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@
44

55
from polyaxon._compiler.contexts.base import BaseContextsManager
66
from polyaxon._connections import V1Connection
7-
from polyaxon._flow import V1CompiledOperation, V1DaskJob, V1DaskReplica
7+
from polyaxon._flow import V1CompiledOperation, V1DaskCluster, V1DaskReplica
88

99

10-
class DaskJobContextsManager(BaseContextsManager):
10+
class DaskClusterContextsManager(BaseContextsManager):
1111
@classmethod
1212
def resolve(
1313
cls,
@@ -21,7 +21,7 @@ def resolve(
2121
) -> Dict:
2222
contexts["init"] = {}
2323
contexts["connections"] = {}
24-
job: V1DaskJob = compiled_operation.run
24+
cluster: V1DaskCluster = compiled_operation.run
2525

2626
def _get_replica(replica: Optional[V1DaskReplica]) -> Dict:
2727
if not replica:
@@ -34,7 +34,6 @@ def _get_replica(replica: Optional[V1DaskReplica]) -> Dict:
3434
)
3535

3636
return {
37-
"job": _get_replica(job.job),
38-
"worker": _get_replica(job.worker),
39-
"scheduler": _get_replica(job.scheduler),
37+
"worker": _get_replica(cluster.worker),
38+
"scheduler": _get_replica(cluster.scheduler),
4039
}

cli/polyaxon/_compiler/contexts/ray_job.py renamed to cli/polyaxon/_compiler/contexts/ray_cluster.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@
44

55
from polyaxon._compiler.contexts.base import BaseContextsManager
66
from polyaxon._connections import V1Connection
7-
from polyaxon._flow import V1CompiledOperation, V1RayJob, V1RayReplica
7+
from polyaxon._flow import V1CompiledOperation, V1RayCluster, V1RayReplica
88

99

10-
class RayJobContextsManager(BaseContextsManager):
10+
class RayClusterContextsManager(BaseContextsManager):
1111
@classmethod
1212
def resolve(
1313
cls,
@@ -21,7 +21,7 @@ def resolve(
2121
) -> Dict:
2222
contexts["init"] = {}
2323
contexts["connections"] = {}
24-
job: V1RayJob = compiled_operation.run
24+
cluster: V1RayCluster = compiled_operation.run
2525

2626
def _get_replica(replica: Optional[V1RayReplica]) -> Dict:
2727
if not replica:
@@ -34,8 +34,10 @@ def _get_replica(replica: Optional[V1RayReplica]) -> Dict:
3434
)
3535

3636
data = {
37-
"head": _get_replica(job.head),
37+
"head": _get_replica(cluster.head),
3838
}
39-
if job.workers:
40-
data["workers"] = {wn: _get_replica(job.workers[wn]) for wn in job.workers}
39+
if cluster.workers:
40+
data["workers"] = {
41+
wn: _get_replica(cluster.workers[wn]) for wn in cluster.workers
42+
}
4143
return data

‎cli/polyaxon/_compiler/resolver/runtime.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ class BaseResolver:
2222
V1RunKind.TUNER,
2323
V1RunKind.WATCHDOG,
2424
V1RunKind.DAG,
25-
V1RunKind.DASKJOB,
26-
V1RunKind.RAYJOB,
25+
V1RunKind.DASKCLUSTER,
26+
V1RunKind.RAYCLUSTER,
2727
V1RunKind.TUNER,
2828
}
2929

‎cli/polyaxon/_constants/metadata.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
META_PORTS = "ports"
77
META_HAS_EARLY_STOPPING = "has_early_stopping"
88
META_HAS_JOBS = "has_jobs"
9+
META_HAS_CLUSTERS = "has_clusters"
910
META_HAS_SERVICES = "has_services"
1011
META_HAS_DAGS = "has_dags"
1112
META_HAS_MATRICES = "has_matrices"

‎cli/polyaxon/_deploy/schemas/operators.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,5 +7,5 @@ class OperatorsConfig(BaseSchemaModel):
77
tfjob: Optional[bool] = None
88
pytorchjob: Optional[bool] = None
99
mpijob: Optional[bool] = None
10-
daskjob: Optional[bool] = None
11-
rayjob: Optional[bool] = None
10+
daskcluster: Optional[bool] = None
11+
raycluster: Optional[bool] = None

‎cli/polyaxon/_env_vars/keys.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,8 @@
102102
ENV_KEYS_TFJOB_ENABLED = "POLYAXON_TFJOB_ENABLED"
103103
ENV_KEYS_PYTORCH_JOB_ENABLED = "POLYAXON_PYTORCH_JOB_ENABLED"
104104
ENV_KEYS_MPIJOB_ENABLED = "POLYAXON_MPIJOB_ENABLED"
105-
ENV_KEYS_RAYJOB_ENABLED = "POLYAXON_RAYJOB_ENABLED"
106-
ENV_KEYS_DASKJOB_ENABLED = "POLYAXON_DASKJOB_ENABLED"
105+
ENV_KEYS_RAYCLUSTER_ENABLED = "POLYAXON_RAYCLUSTER_ENABLED"
106+
ENV_KEYS_DASKCLUSTER_ENABLED = "POLYAXON_DASKCLUSTER_ENABLED"
107107

108108
# Sandbox
109109
ENV_KEYS_SANDBOX_PORT = "POLYAXON_SANDBOX_PORT"

‎cli/polyaxon/_flow/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@
6969
V1CleanPodPolicy,
7070
V1CloningKind,
7171
V1Dag,
72-
V1DaskJob,
72+
V1DaskCluster,
7373
V1DaskReplica,
7474
V1Job,
7575
V1KFReplica,
@@ -78,7 +78,7 @@
7878
V1PipelineKind,
7979
V1PytorchElasticPolicy,
8080
V1PytorchJob,
81-
V1RayJob,
81+
V1RayCluster,
8282
V1RayReplica,
8383
V1RunEdgeKind,
8484
V1RunKind,

‎cli/polyaxon/_flow/component/component.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ class V1Component(
4747
hooks: List[[V1Hook](/docs/automation/hooks/)], optional
4848
inputs: [V1IO](/docs/core/specification/io/), optional
4949
outputs: [V1IO](/docs/core/specification/io/), optional
50-
run: Union[[V1Job](/docs/experimentation/jobs/), [V1Service](/docs/experimentation/services/), [V1TFJob](/docs/experimentation/distributed/tf-jobs/), [V1PytorchJob](/docs/experimentation/distributed/pytorch-jobs/), [V1MPIJob](/docs/experimentation/distributed/mpi-jobs/), [V1RayJob](/docs/experimentation/distributed/ray-jobs/), [V1DaskJob](/docs/experimentation/distributed/dask-jobs/), [V1Dag](/docs/automation/flow-engine/specification/)] # noqa
50+
run: Union[[V1Job](/docs/experimentation/jobs/), [V1Service](/docs/experimentation/services/), [V1TFJob](/docs/experimentation/distributed/tf-jobs/), [V1PytorchJob](/docs/experimentation/distributed/pytorch-jobs/), [V1MPIJob](/docs/experimentation/distributed/mpi-jobs/), [V1RayCluster](/docs/experimentation/distributed/ray-jobs/), [V1DaskCluster](/docs/experimentation/distributed/dask-jobs/), [V1Dag](/docs/automation/flow-engine/specification/)] # noqa
5151
template: [V1Template](/docs/core/specification/template/), optional
5252
5353
## YAML usage
@@ -346,8 +346,8 @@ class V1Component(
346346
* [V1PytorchJob](/docs/experimentation/distributed/pytorch-jobs/): for running distributed
347347
Pytorch training job.
348348
* [V1MPIJob](/docs/experimentation/distributed/mpi-jobs/): for running distributed MPI job.
349-
* [V1RayJob](/docs/experimentation/distributed/ray-jobs/): for running a ray job.
350-
* [V1DaskJob](/docs/experimentation/distributed/dask-jobs/): for running a Dask job.
349+
* [V1RayCluster](/docs/experimentation/distributed/ray-jobs/): for running a ray job.
350+
* [V1DaskCluster](/docs/experimentation/distributed/dask-jobs/): for running a Dask job.
351351
* [V1Dag](/docs/automation/flow-engine/specification/): for running a DAG/workflow.
352352
353353
### isApproved

0 commit comments

Comments
 (0)