remove: ddp

Author: VanillaSpoon

src/codeflare_sdk/__init__.py

@@ -14,6 +14,6 @@
     get_cluster,
 )
 
-from .job import JobDefinition, Job, DDPJobDefinition, DDPJob, RayJobClient
+from .job import JobDefinition, Job, RayJobClient
 
 from .utils import generate_cert

src/codeflare_sdk/job/__init__.py

@@ -1,3 +1,3 @@
-from .jobs import JobDefinition, Job, DDPJobDefinition, DDPJob
+from .jobs import JobDefinition, Job
 
 from .ray_jobs import RayJobClient

src/codeflare_sdk/job/jobs.py

@@ -12,15 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-The jobs sub-module contains methods needed to submit Distributed Data Parallel(DDP) jobs to Ray Clusters created by the CodeFlare SDK.
-"""
-
 import abc
 from typing import TYPE_CHECKING, Optional, Dict, List
 from pathlib import Path
 
-from torchx.components.dist import ddp
 from torchx.runner import get_runner, Runner
 from torchx.schedulers.ray_scheduler import RayScheduler
 from torchx.specs import AppHandle, parse_app_handle, AppDryRunInfo
@@ -47,161 +42,3 @@ def status(self):
 
     def logs(self):
         pass
-
-
-class DDPJobDefinition(JobDefinition):
-    def __init__(
-        self,
-        script: Optional[str] = None,
-        m: Optional[str] = None,
-        script_args: Optional[List[str]] = None,
-        name: Optional[str] = None,
-        cpu: Optional[int] = None,
-        gpu: Optional[int] = None,
-        memMB: Optional[int] = None,
-        h: Optional[str] = None,
-        j: Optional[str] = None,
-        env: Optional[Dict[str, str]] = None,
-        max_retries: int = 0,
-        mounts: Optional[List[str]] = None,
-        rdzv_port: int = 29500,
-        rdzv_backend: str = None,
-        scheduler_args: Optional[Dict[str, str]] = None,
-        image: Optional[str] = None,
-        workspace: Optional[str] = f"file://{Path.cwd()}",
-    ):
-        if bool(script) == bool(m):  # logical XOR
-            raise ValueError(
-                "Exactly one of the following arguments must be defined: [script, m]."
-            )
-        self.script = script
-        self.m = m
-        self.script_args: List[str] = script_args if script_args is not None else []
-        self.name = name
-        self.cpu = cpu
-        self.gpu = gpu
-        self.memMB = memMB
-        self.h = h
-        self.j = j
-        self.env: Dict[str, str] = env if env is not None else dict()
-        self.max_retries = max_retries
-        self.mounts: List[str] = mounts if mounts is not None else []
-        self.rdzv_port = rdzv_port
-        self.rdzv_backend = rdzv_backend
-        self.scheduler_args: Dict[str, str] = (
-            scheduler_args if scheduler_args is not None else dict()
-        )
-        self.image = image
-        self.workspace = workspace
-
-    def _dry_run(self, cluster: "Cluster"):
-        j = f"{cluster.config.num_workers}x{max(cluster.config.num_gpus, 1)}"  # # of proc. = # of gpus
-        runner = get_runner(ray_client=cluster.job_client)
-        runner._scheduler_instances["ray"] = RayScheduler(
-            session_name=runner._name, ray_client=cluster.job_client
-        )
-        return (
-            runner.dryrun(
-                app=ddp(
-                    *self.script_args,
-                    script=self.script,
-                    m=self.m,
-                    name=self.name,
-                    h=self.h,
-                    cpu=self.cpu if self.cpu is not None else cluster.config.max_cpus,
-                    gpu=self.gpu if self.gpu is not None else cluster.config.num_gpus,
-                    memMB=self.memMB
-                    if self.memMB is not None
-                    else cluster.config.max_memory * 1024,
-                    j=self.j if self.j is not None else j,
-                    env=self.env,
-                    max_retries=self.max_retries,
-                    rdzv_port=self.rdzv_port,
-                    rdzv_backend=self.rdzv_backend
-                    if self.rdzv_backend is not None
-                    else "static",
-                    mounts=self.mounts,
-                ),
-                scheduler=cluster.torchx_scheduler,
-                cfg=cluster.torchx_config(**self.scheduler_args),
-                workspace=self.workspace,
-            ),
-            runner,
-        )
-
-    def _missing_spec(self, spec: str):
-        raise ValueError(f"Job definition missing arg: {spec}")
-
-    def _dry_run_no_cluster(self):
-        if self.scheduler_args is not None:
-            if self.scheduler_args.get("namespace") is None:
-                self.scheduler_args["namespace"] = get_current_namespace()
-        runner = get_runner()
-        return (
-            runner.dryrun(
-                app=ddp(
-                    *self.script_args,
-                    script=self.script,
-                    m=self.m,
-                    name=self.name
-                    if self.name is not None
-                    else self._missing_spec("name"),
-                    h=self.h,
-                    cpu=self.cpu
-                    if self.cpu is not None
-                    else self._missing_spec("cpu (# cpus per worker)"),
-                    gpu=self.gpu
-                    if self.gpu is not None
-                    else self._missing_spec("gpu (# gpus per worker)"),
-                    memMB=self.memMB
-                    if self.memMB is not None
-                    else self._missing_spec("memMB (memory in MB)"),
-                    j=self.j
-                    if self.j is not None
-                    else self._missing_spec(
-                        "j (`workers`x`procs`)"
-                    ),  # # of proc. = # of gpus,
-                    env=self.env,  # should this still exist?
-                    max_retries=self.max_retries,
-                    rdzv_port=self.rdzv_port,  # should this still exist?
-                    rdzv_backend=self.rdzv_backend
-                    if self.rdzv_backend is not None
-                    else "c10d",
-                    mounts=self.mounts,
-                    image=self.image
-                    if self.image is not None
-                    else self._missing_spec("image"),
-                ),
-                scheduler="kubernetes_mcad",
-                cfg=self.scheduler_args,
-                workspace="",
-            ),
-            runner,
-        )
-
-    def submit(self, cluster: "Cluster" = None) -> "Job":
-        return DDPJob(self, cluster)
-
-
-class DDPJob(Job):
-    def __init__(self, job_definition: "DDPJobDefinition", cluster: "Cluster" = None):
-        self.job_definition = job_definition
-        self.cluster = cluster
-        if self.cluster:
-            definition, runner = job_definition._dry_run(cluster)
-            self._app_handle = runner.schedule(definition)
-            self._runner = runner
-        else:
-            definition, runner = job_definition._dry_run_no_cluster()
-            self._app_handle = runner.schedule(definition)
-            self._runner = runner
-        all_jobs.append(self)
-
-    def status(self) -> str:
-        return self._runner.status(self._app_handle)
-
-    def logs(self) -> str:
-        return "".join(self._runner.log_lines(self._app_handle, None))
-
-    def cancel(self):
-        self._runner.cancel(self._app_handle)

tests/e2e/mnist_raycluster_sdk_oauth_test.py

@@ -5,7 +5,7 @@
 from torchx.specs.api import AppState, is_terminal
 
 from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication
-from codeflare_sdk.job.jobs import DDPJobDefinition
+from codeflare_sdk.job.jobs import JobDefinition
 
 import pytest
 
@@ -97,7 +97,7 @@ def assert_jobsubmit_withoutLogin(self, cluster):
 
     def assert_jobsubmit_withlogin(self, cluster):
         self.assert_appwrapper_exists()
-        jobdef = DDPJobDefinition(
+        jobdef = JobDefinition(
             name="mnist",
             script="./tests/e2e/mnist.py",
             scheduler_args={"requirements": "./tests/e2e/mnist_pip_requirements.txt"},

tests/e2e/mnist_raycluster_sdk_test.py

@@ -10,7 +10,7 @@
 from torchx.specs.api import AppState, is_terminal
 
 from codeflare_sdk.cluster.cluster import Cluster, ClusterConfiguration
-from codeflare_sdk.job.jobs import DDPJobDefinition
+from codeflare_sdk.job.jobs import JobDefinition
 
 import pytest
 
@@ -69,6 +69,7 @@ def run_mnist_raycluster_sdk(self):
                 num_gpus=0,
                 instascale=False,
                 image=ray_image,
+                ingress_domain="apps.cluster.awsroute.org",
                 ingress_options=ingress_options,
                 write_to_file=True,
             )
@@ -86,7 +87,7 @@ def run_mnist_raycluster_sdk(self):
 
         cluster.details()
 
-        jobdef = DDPJobDefinition(
+        jobdef = JobDefinition(
             name="mnist",
             script="./tests/e2e/mnist.py",
             scheduler_args={"requirements": "./tests/e2e/mnist_pip_requirements.txt"},

tests/e2e/mnist_rayjob.py

@@ -5,15 +5,15 @@
 from torchx.specs.api import AppState, is_terminal
 
 from codeflare_sdk.cluster.cluster import get_cluster
-from codeflare_sdk.job.jobs import DDPJobDefinition
+from codeflare_sdk.job.jobs import JobDefinition
 
 namespace = sys.argv[1]
 
 cluster = get_cluster("mnist", namespace)
 
 cluster.details()
 
-jobdef = DDPJobDefinition(
+jobdef = JobDefinition(
     name="mnist",
     script="mnist.py",
     scheduler_args={"requirements": "requirements.txt"},