diff --git a/python/lightning_sdk/api/job_api.py b/python/lightning_sdk/api/job_api.py index 12c0b7ea..b0c4da84 100644 --- a/python/lightning_sdk/api/job_api.py +++ b/python/lightning_sdk/api/job_api.py @@ -143,6 +143,7 @@ def submit_job( reuse_snapshot: bool = True, scratch_disks: Optional[Dict[str, int]] = None, placement_group_id: Optional[str] = None, + num_machines: int = 1, ) -> V1Job: """Submit a v2 job and return the created job object. @@ -164,10 +165,13 @@ def submit_job( reuse_snapshot: Whether to reuse the Studio's existing filesystem snapshot. scratch_disks: Optional mapping of scratch-disk mount paths to their sizes in GiB. placement_group_id: Optional placement group identifier for colocating the job. + num_machines: Must be 1 for single-machine jobs. Kept for parity with ``MMTApiV2.submit_job``. Returns: The newly created ``V1Job`` object. """ + if num_machines != 1: + raise ValueError("JobApiV2 only supports single-machine jobs (num_machines=1)") if scratch_disks is not None: sanitized_scratch_disks = {} for k, v in scratch_disks.items(): @@ -342,7 +346,7 @@ def stop_job(self, job_id: str, teamspace_id: str) -> None: break time.sleep(1) - def delete_job(self, job_id: str, teamspace_id: str, cloudspace_id: Optional[str]) -> None: + def delete_job(self, job_id: str, teamspace_id: str, cloudspace_id: Optional[str] = None) -> None: """Permanently delete a v2 job. Args: @@ -622,3 +626,14 @@ def get_total_cost(self, job: V1Job) -> float: The total cost incurred by the job, expressed in US dollars. """ return job.total_cost + + def get_num_machines(self, job: V1Job) -> int: + """Return the number of machines for this job. + + Single-machine jobs always report 1. Kept for parity with ``MMTApiV2.get_num_machines``. + """ + return 1 + + def list_mmt_subjobs(self, job_id: str, teamspace_id: str) -> List[V1Job]: + """Single-machine jobs have no sub-jobs. Kept for parity with ``MMTApiV2.list_mmt_subjobs``.""" + return [] diff --git a/python/lightning_sdk/api/mmt_api.py b/python/lightning_sdk/api/mmt_api.py index ad6c4410..1c0d934a 100644 --- a/python/lightning_sdk/api/mmt_api.py +++ b/python/lightning_sdk/api/mmt_api.py @@ -51,6 +51,7 @@ def submit_job( max_runtime: Optional[int], reuse_snapshot: bool, placement_group_id: Optional[str] = None, + scratch_disks: Optional[Dict[str, int]] = None, ) -> V1MultiMachineJob: """Submit a v2 multi-machine job and return the created job object. @@ -72,10 +73,13 @@ def submit_job( max_runtime: Maximum allowed runtime in seconds, or ``None`` for no limit. reuse_snapshot: Whether to reuse the Studio's existing filesystem snapshot. placement_group_id: Optional placement group identifier for colocating the job. + scratch_disks: Not supported for multi-machine jobs. Kept for parity with ``JobApiV2.submit_job``. Returns: The newly created ``V1MultiMachineJob`` object. """ + if scratch_disks: + raise ValueError("scratch_disks are not supported for multi-machine jobs") body = self._create_mmt_body( name=name, num_machines=num_machines, @@ -242,12 +246,13 @@ def stop_job(self, job_id: str, teamspace_id: str) -> None: break time.sleep(1) - def delete_job(self, job_id: str, teamspace_id: str) -> None: + def delete_job(self, job_id: str, teamspace_id: str, cloudspace_id: Optional[str] = None) -> None: """Permanently delete a multi-machine job. Args: job_id: The unique identifier of the multi-machine job to delete. teamspace_id: The ID of the teamspace that owns the job. + cloudspace_id: Ignored. Kept for parity with ``JobApiV2.delete_job``. """ self._client.jobs_service_delete_multi_machine_job(project_id=teamspace_id, id=job_id) diff --git a/python/lightning_sdk/cli/job/inspect.py b/python/lightning_sdk/cli/job/inspect.py index 856ca024..368c24a1 100644 --- a/python/lightning_sdk/cli/job/inspect.py +++ b/python/lightning_sdk/cli/job/inspect.py @@ -6,7 +6,7 @@ from rich.console import Console from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace @click.command("inspect", cls=LightningCommand) @@ -20,9 +20,19 @@ "If not specified, uses the configured default teamspace." ), ) +@click.option("--rank", type=int, default=None, help="Inspect one machine in a multi-machine job.") @click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON (inspect always emits JSON).") -def inspect_job(name: Optional[str] = None, teamspace: Optional[str] = None, as_json: bool = False) -> None: +def inspect_job( + name: Optional[str] = None, + teamspace: Optional[str] = None, + rank: Optional[int] = None, + as_json: bool = False, +) -> None: """Inspect a job for further details as JSON.""" resolved_teamspace = resolve_teamspace(teamspace) job = resolve_job(name, resolved_teamspace) + if job.is_multi_machine is True and rank is not None: + job = resolve_job_machine(job, rank) + elif rank is not None: + raise click.UsageError("--rank is only supported for multi-machine jobs.") Console().print(job.json()) diff --git a/python/lightning_sdk/cli/job/list.py b/python/lightning_sdk/cli/job/list.py index 29addc11..3db95d5c 100644 --- a/python/lightning_sdk/cli/job/list.py +++ b/python/lightning_sdk/cli/job/list.py @@ -1,10 +1,16 @@ """Job list command.""" +from contextlib import suppress from typing import Optional import rich_click as click +from rich.console import Console +from rich.table import Table +from lightning_sdk.cli.utils.json_output import echo_json from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.cli.utils.resource_resolution import resolve_teamspace +from lightning_sdk.models import _list_teamspaces @click.command("list", cls=LightningCommand) @@ -38,7 +44,55 @@ def list_jobs( sort_by: Optional[str] = None, as_json: bool = False, ) -> None: - """List jobs for a given teamspace.""" - from lightning_sdk.cli.legacy.list import jobs + """List jobs for a given teamspace. - jobs.callback(teamspace=teamspace, all=all, sort_by=sort_by, as_json=as_json) + Includes both single- and multi-machine jobs. + """ + resources = [] + if all and not teamspace: + for teamspace_slug in _list_teamspaces(): + resolved = resolve_teamspace(teamspace_slug) + resources.extend(resolved.jobs) + else: + resolved = resolve_teamspace(teamspace) + resources.extend(resolved.jobs) + + rows = [] + for job in resources: + job._prevent_refetch_latest = True + with suppress(RuntimeError): + rows.append( + { + "name": job.name, + "teamspace": f"{job.teamspace.owner.name}/{job.teamspace.name}", + "studio": job.studio_name, + "image": job.image, + "status": str(job.status) if job.status is not None else None, + "machine": str(job.machine), + "num_machines": getattr(job, "num_machines", 1), + "total_cost": round(job.total_cost, 3), + "_cloud_account": str(getattr(job, "cloud_account", "") or ""), + } + ) + + sort_key = "_cloud_account" if sort_by == "cloud-account" else sort_by or "name" + rows.sort(key=lambda row: str(row.get(sort_key) or "")) + if as_json: + echo_json([{key: value for key, value in row.items() if not key.startswith("_")} for row in rows]) + return + + table = Table(pad_edge=True) + for column in ("Name", "Teamspace", "Studio", "Image", "Status", "Machine", "Num Machines", "Total Cost"): + table.add_column(column) + for row in rows: + table.add_row( + row["name"], + row["teamspace"], + row["studio"], + row["image"], + row["status"], + row["machine"], + str(row["num_machines"]), + f"{row['total_cost']:.3f}", + ) + Console().print(table) diff --git a/python/lightning_sdk/cli/job/logs.py b/python/lightning_sdk/cli/job/logs.py index fd9e51d2..530fd202 100644 --- a/python/lightning_sdk/cli/job/logs.py +++ b/python/lightning_sdk/cli/job/logs.py @@ -1,5 +1,6 @@ """Job logs command.""" +from contextlib import suppress from typing import Optional import rich_click as click @@ -7,7 +8,7 @@ from lightning_sdk.api.logs_api import SEVERITIES from lightning_sdk.cli.utils.logging import LightningCommand from lightning_sdk.cli.utils.logs import LogSelection, read_logs, resolve_time -from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace @click.command("logs", cls=LightningCommand) @@ -19,7 +20,7 @@ ) @click.option("--follow", "-f", is_flag=True, default=False, help="Stream new log lines as they are produced.") @click.option("--tail", type=int, default=None, help="Only show the last N lines.") -@click.option("--rank", type=int, default=None, help="Distributed job rank to read from (running jobs only).") +@click.option("--rank", type=int, default=None, help="Machine rank to read from in a multi-machine job.") @click.option("--timestamps", is_flag=True, default=False, help="Prepend each line with its ISO-8601 timestamp.") @click.option("--since", default=None, help='Only include lines at or after this time (e.g. "2h", RFC3339).') @click.option("--until", default=None, help='Only include lines at or before this time (e.g. "30m", RFC3339).') @@ -49,15 +50,31 @@ def logs_job( Prints a snapshot of the logs available so far. Pass --follow to stream new lines from a running job until it finishes or you press Ctrl-C. --query and --severity are applied by the server, to both the snapshot and the stream. + Multi-machine logs are merged unless --rank selects one machine, which opens + that machine's per-job websocket (same path as a single-machine job with + --rank). """ resolved_teamspace = resolve_teamspace(teamspace) - job = resolve_job(name, resolved_teamspace) + resource = resolve_job(name, resolved_teamspace) + selected_rank = resource.is_multi_machine is True and rank is not None + job = resolve_job_machine(resource, rank) if selected_rank else resource if as_json: - if rank is not None: + if rank is not None and not selected_rank: raise click.ClickException("--rank is not supported with --json.") + if job.is_multi_machine is True: + labels: dict = {} + with suppress(Exception): + labels = {machine.resource_id: machine.name for machine in job.machines} + selection = LogSelection( + teamspace_id=resolved_teamspace.id, + mmt_id=job.resource_id, + labels=labels, + ) + else: + selection = LogSelection(teamspace_id=resolved_teamspace.id, job_ids=[job.resource_id]) read_logs( - LogSelection(teamspace_id=resolved_teamspace.id, job_ids=[job.resource_id]), + selection, query=query, severity=severity, since=resolve_time(since, "--since"), @@ -69,16 +86,20 @@ def logs_job( return try: - logs = job.logs( - follow=follow, - tail=tail, - rank=rank, - timestamps=timestamps, - since=resolve_time(since, "--since"), - until=resolve_time(until, "--until"), - query=query, - severity=severity, - ) + log_options = { + "follow": follow, + "tail": tail, + "timestamps": timestamps, + "since": resolve_time(since, "--since"), + "until": resolve_time(until, "--until"), + "query": query, + "severity": severity, + } + if job.is_multi_machine is not True: + # Any non-None rank routes Job through the legacy per-job websocket (server-side + # tail). For a selected MMT machine the process rank on that node is 0. + log_options["rank"] = 0 if selected_rank else rank + logs = job.logs(**log_options) if follow: for line in logs: click.echo(line) diff --git a/python/lightning_sdk/cli/job/run.py b/python/lightning_sdk/cli/job/run.py index 8e69ff88..1d86a7aa 100644 --- a/python/lightning_sdk/cli/job/run.py +++ b/python/lightning_sdk/cli/job/run.py @@ -18,6 +18,14 @@ @click.command("run", cls=LightningCommand) @click.option("--name", default=None, help="The name of the job. Needs to be unique within the teamspace.") +@click.option( + "--num-machines", + "--num_machines", + default=1, + show_default=True, + type=click.IntRange(min=1), + help="The number of machines to run on.", +) @click.option( "--machine", default="CPU", @@ -119,6 +127,7 @@ @click.option("--json", "as_json", is_flag=True, default=False, help="Output the created job as JSON.") def run_job( name: Optional[str] = None, + num_machines: int = 1, machine: str = "CPU", command: Optional[str] = None, studio: Optional[str] = None, @@ -136,7 +145,10 @@ def run_job( path_mappings: str = "", as_json: bool = False, ) -> None: - """Run async workloads using a docker image or studio.""" + """Run async workloads using a docker image or studio. + + Pass --num-machines greater than 1 to run a multi-machine job. + """ if not name: from datetime import datetime @@ -159,23 +171,25 @@ def run_job( for value in env: env_dict.update(_resolve_envs(value)) - job = Job.run( - name=name, - machine=machine_enum, - command=command, - studio=studio, - image=image, - teamspace=resolved_teamspace, - org=org, - user=user, - cloud=cloud, - env=env_dict, - interruptible=interruptible, - image_credentials=image_credentials, - cloud_account_auth=cloud_account_auth, - entrypoint=entrypoint, - path_mappings=path_mappings_dict, - ) + run_kwargs = { + "name": name, + "machine": machine_enum, + "command": command, + "studio": studio, + "image": image, + "teamspace": resolved_teamspace, + "org": org, + "user": user, + "cloud": cloud, + "env": env_dict, + "interruptible": interruptible, + "image_credentials": image_credentials, + "cloud_account_auth": cloud_account_auth, + "entrypoint": entrypoint, + "path_mappings": path_mappings_dict, + "num_machines": num_machines, + } + job = Job.run(**run_kwargs) if as_json: echo_json( diff --git a/python/lightning_sdk/cli/job/ssh.py b/python/lightning_sdk/cli/job/ssh.py index b3efe26e..ae39ee22 100644 --- a/python/lightning_sdk/cli/job/ssh.py +++ b/python/lightning_sdk/cli/job/ssh.py @@ -6,7 +6,7 @@ import rich_click as click from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal from lightning_sdk.status import Status @@ -31,27 +31,34 @@ def _ssh_user_for_job_id(job_id: str) -> str: "If not specified, uses the configured default teamspace." ), ) +@click.option("--rank", type=int, default=None, help="Machine rank for a multi-machine job. Defaults to 0.") def ssh_job( name: str, teamspace: Optional[str] = None, + rank: Optional[int] = None, ) -> None: """SSH into a running job. Example: lightning job ssh my-job """ - ssh_impl(name=name, teamspace=teamspace) + ssh_impl(name=name, teamspace=teamspace, rank=rank) def ssh_impl( name: Optional[str], teamspace: Optional[str], + rank: Optional[int] = None, ) -> None: if not name: raise click.UsageError("Missing job name. Pass NAME.") resolved_teamspace = resolve_teamspace(teamspace) job = resolve_job(name, resolved_teamspace) + if job.is_multi_machine is True: + job = resolve_job_machine(job, rank if rank is not None else 0) + elif rank is not None: + raise click.UsageError("--rank is only supported for multi-machine jobs.") if job.status != Status.Running: raise click.ClickException( diff --git a/python/lightning_sdk/cli/mmt/ssh.py b/python/lightning_sdk/cli/mmt/ssh.py index 747d2491..ad5cd547 100644 --- a/python/lightning_sdk/cli/mmt/ssh.py +++ b/python/lightning_sdk/cli/mmt/ssh.py @@ -6,10 +6,8 @@ import rich_click as click from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.cli.utils.resource_resolution import resolve_mmt, resolve_teamspace +from lightning_sdk.cli.utils.resource_resolution import resolve_job_machine, resolve_mmt, resolve_teamspace from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal -from lightning_sdk.job import Job -from lightning_sdk.mmt import MMT from lightning_sdk.status import Status _SSH_HOST = "ssh.lightning.ai" @@ -22,30 +20,6 @@ def _ssh_user_for_job_id(job_id: str) -> str: return f"j_{suffix}" -def _machine_for_rank(mmt: MMT, rank: int) -> Job: - machines = mmt.machines - if not machines: - raise click.ClickException(f"Multi-machine job '{mmt.name}' has no machines to SSH into.") - - expected = f"{mmt.name}-{rank}" - for machine in machines: - if machine.name == expected: - return machine - - prefix = f"{mmt.name}-" - available_ranks = [] - for machine in machines: - if not machine.name.startswith(prefix): - continue - suffix = machine.name[len(prefix) :] - if suffix.isdigit(): - available_ranks.append(int(suffix)) - available = ", ".join(str(r) for r in sorted(available_ranks)) - raise click.ClickException( - f"Rank {rank} not found on multi-machine job '{mmt.name}'. Available ranks: {available or 'none'}." - ) - - @click.command("ssh", cls=LightningCommand) @click.argument("name") @click.option( @@ -84,7 +58,7 @@ def ssh_impl( ) -> None: resolved_teamspace = resolve_teamspace(teamspace) mmt = resolve_mmt(name, resolved_teamspace) - job = _machine_for_rank(mmt, rank) + job = resolve_job_machine(mmt, rank) if job.status != Status.Running: raise click.ClickException( diff --git a/python/lightning_sdk/cli/utils/resource_resolution.py b/python/lightning_sdk/cli/utils/resource_resolution.py index b0331fdf..7d05c384 100644 --- a/python/lightning_sdk/cli/utils/resource_resolution.py +++ b/python/lightning_sdk/cli/utils/resource_resolution.py @@ -78,6 +78,36 @@ def resolve_job(name: Optional[str], teamspace: Teamspace) -> Job: raise click.UsageError(f"Could not resolve job '{name}' in teamspace '{teamspace.name}'.") from ex +def resolve_job_machine(job: Job, rank: int) -> Job: + """Resolve one machine in a multi-machine job by rank.""" + machines = job.machines + if not machines: + raise click.ClickException(f"Job '{job.name}' has no machines.") + + for machine in machines: + if machine.rank == rank: + return machine + + # Fallback for older naming when rank is missing on the machine object. + expected = f"{job.name}-{rank}" + for machine in machines: + if machine.name == expected: + return machine + + available_ranks: set[int] = set() + prefix = f"{job.name}-" + for machine in machines: + if machine.rank is not None: + available_ranks.add(machine.rank) + continue + if machine.name.startswith(prefix): + suffix = machine.name[len(prefix) :] + if suffix.isdigit(): + available_ranks.add(int(suffix)) + available = ", ".join(str(value) for value in sorted(available_ranks)) + raise click.ClickException(f"Rank {rank} not found on job '{job.name}'. Available ranks: {available or 'none'}.") + + def resolve_mmt(name: Optional[str], teamspace: Teamspace) -> MMT: if not name: raise click.UsageError("Missing multi-machine job name. Pass JOB.") diff --git a/python/lightning_sdk/job.py b/python/lightning_sdk/job.py index 65f9de35..f6a5a5b7 100644 --- a/python/lightning_sdk/job.py +++ b/python/lightning_sdk/job.py @@ -1,10 +1,11 @@ import warnings from pathlib import PurePath -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, TypedDict, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Tuple, TypedDict, Union from lightning_sdk.api.cloud_account_api import CloudAccountApi from lightning_sdk.api.job_api import JobApiV2 from lightning_sdk.api.logs_api import LogsApi +from lightning_sdk.api.mmt_api import MMTApiV2 from lightning_sdk.api.utils import AccessibleResource, _get_cloud_url, raise_access_error_if_not_allowed from lightning_sdk.status import Status from lightning_sdk.utils.logging import TrackCallsMeta @@ -121,7 +122,7 @@ class JobDict(TypedDict): class Job(metaclass=TrackCallsMeta): - """Submit and manage single-machine jobs on the Lightning AI Platform.""" + """Submit and manage jobs on the Lightning AI Platform.""" def __init__( self, @@ -131,6 +132,7 @@ def __init__( user: Union[str, "User", None] = None, *, _fetch_job: bool = True, + _num_machines: int = 1, ) -> None: """Fetch already existing jobs. @@ -148,6 +150,9 @@ def __init__( when ``_fetch_job=True``. PermissionError: If the user does not have access to jobs in the given teamspace. """ + if _num_machines < 1: + raise ValueError("A job needs to run on at least one machine") + teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user) if teamspace is None: raise ValueError( @@ -162,7 +167,9 @@ def __init__( self._job = None self._prevent_refetch_latest = False self._cloud_account_api = CloudAccountApi() - self._job_api = JobApiV2() + self._standalone_job_api = JobApiV2() + self._mmt_job_api = MMTApiV2() + self._num_machines = _num_machines self._logs_api = LogsApi() if _fetch_job: @@ -175,6 +182,21 @@ def __init__( raise ValueError(f"Job {name} does not exist in Teamspace {teamspace.name}") from None raise + @property + def _job_api(self) -> Union[JobApiV2, MMTApiV2]: + return self._mmt_job_api if self._num_machines > 1 else self._standalone_job_api + + def _attach_job(self, job: Any) -> None: + """Bind a fetched job payload and sync ``num_machines`` from it.""" + from lightning_sdk.lightning_cloud.openapi import V1Job, V1MultiMachineJob + + self._job = job + # Only a known payload type changes the machine count; anything else keeps the current one. + if isinstance(job, V1MultiMachineJob): + self._num_machines = job.machines if job.machines and job.machines > 1 else max(self._num_machines, 2) + elif isinstance(job, V1Job): + self._num_machines = 1 + @classmethod def run( cls, @@ -197,12 +219,14 @@ def run( reuse_snapshot: bool = True, scratch_disks: Optional[Dict[str, int]] = None, placement_group_id: Optional[str] = None, + num_machines: int = 1, ) -> "Job": """Run async workloads using a docker image or a compute environment from your studio. Args: name: The name of the job. Needs to be unique within the teamspace. machine: The machine type to run the job on. + num_machines: The number of machines to run on. Defaults to one. command: The command to run inside your job. Required if using a studio. Optional if using an image. If not provided for images, will run the container entrypoint and default command. studio: The studio env to run the job with. Mutually exclusive with image. @@ -248,6 +272,10 @@ def run( if not name: raise ValueError("A job needs to have a name!") + if num_machines < 1: + raise ValueError("A job needs to run on at least one machine") + if num_machines > 1 and scratch_disks: + raise ValueError("scratch_disks are not supported for multi-machine jobs") if image is None: if not isinstance(studio, Studio): @@ -307,10 +335,11 @@ def run( elif entrypoint == "" or entrypoint is None: entrypoint = None - job = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False) + job = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False, _num_machines=num_machines) submit_cloud = cloud if cloud_account is None else None job._submit( + num_machines=num_machines, machine=machine, cloud=submit_cloud, command=command, @@ -350,7 +379,13 @@ def _submit( reuse_snapshot: bool = True, scratch_disks: Optional[Dict[str, int]] = None, placement_group_id: Optional[str] = None, + num_machines: int = 1, ) -> "Job": + if num_machines < 1: + raise ValueError("A job needs to run on at least one machine") + if num_machines > 1 and scratch_disks: + raise ValueError("scratch_disks are not supported for multi-machine jobs") + if studio is not None: studio_id = studio._studio.id if image is not None: @@ -391,6 +426,7 @@ def _submit( if ".." in path.parts: raise ValueError("scratch_disk path cannot contain '..'") + self._num_machines = num_machines submitted = self._job_api.submit_job( name=self.name, command=command, @@ -407,17 +443,18 @@ def _submit( path_mappings=path_mappings, max_runtime=max_runtime, reuse_snapshot=reuse_snapshot, - scratch_disks=scratch_disks, placement_group_id=placement_group_id, + num_machines=num_machines, + scratch_disks=scratch_disks, ) - if submitted.name != self._name: + if num_machines <= 1 and submitted.name != self._name: warnings.warn( f"Job name '{self._name}' was already taken in this teamspace; " f"the job was created as '{submitted.name}' instead.", stacklevel=2, ) - self._job = submitted + self._attach_job(submitted) self._name = submitted.name return self @@ -428,10 +465,11 @@ def stop(self) -> None: self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id) def delete(self) -> None: + cloudspace_id = None if self.is_multi_machine else self._guaranteed_job.spec.cloudspace_id self._job_api.delete_job( job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id, - cloudspace_id=self._guaranteed_job.spec.cloudspace_id, + cloudspace_id=cloudspace_id, ) def wait(self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False) -> None: @@ -485,6 +523,8 @@ def machine(self) -> Union["Machine", str]: @property def public_ip(self) -> Optional[str]: + if self.is_multi_machine: + return None try: return self._job.public_ip_address except AttributeError: @@ -501,6 +541,8 @@ def resource_id(self) -> Optional[str]: @property def private_ip_address(self) -> Optional[str]: + if self.is_multi_machine: + return None return self._guaranteed_job.private_ip_address @property @@ -509,10 +551,41 @@ def placement_group_id(self) -> Optional[str]: @property def rank(self) -> Optional[int]: + if self.is_multi_machine: + return None return self._guaranteed_job.spec.rank + @property + def is_multi_machine(self) -> bool: + """Whether this object represents a multi-machine parent job.""" + return self._num_machines > 1 + + @property + def num_machines(self) -> int: + """The number of machines allocated to this job.""" + return self._num_machines + + @property + def machines(self) -> Tuple["Job", ...]: + """The rank-ordered machines in this job.""" + if not self.is_multi_machine: + return (self,) + + subjobs = sorted( + self._job_api.list_mmt_subjobs(self._guaranteed_job.id, self.teamspace.id), + key=lambda job: job.spec.rank, + ) + machines = [] + for subjob in subjobs: + job = Job(name=subjob.name, teamspace=self.teamspace, _fetch_job=False, _num_machines=1) + job._attach_job(subjob) + machines.append(job) + return tuple(machines) + @property def artifact_path(self) -> Optional[str]: + if self.is_multi_machine: + raise NotImplementedError if self._guaranteed_job.spec.image != "": if self._guaranteed_job.spec.artifacts_destination: ( @@ -527,12 +600,16 @@ def artifact_path(self) -> Optional[str]: @property def snapshot_path(self) -> Optional[str]: + if self.is_multi_machine: + raise NotImplementedError if self._guaranteed_job.spec.image != "": return None return f"/teamspace/jobs/{self._guaranteed_job.name}/snapshot" @property def share_path(self) -> Optional[str]: + if self.is_multi_machine: + return None raise NotImplementedError("Not implemented yet") @property @@ -579,6 +656,18 @@ def _compute_logs( severity: Optional[str] = None, ) -> Union[str, Iterator[str]]: """Fetch the logs, dispatching on job state. See :attr:`logs` for the public API.""" + if self.is_multi_machine: + return self._compute_multi_machine_logs( + follow=follow, + tail=tail, + rank=rank, + timestamps=timestamps, + since=since, + until=until, + query=query, + severity=severity, + ) + status = self.status if rank is not None: @@ -609,6 +698,72 @@ def _compute_logs( # keeping the return type consistent with the live-follow path. return iter(lines) if follow else "\n".join(lines) + def _compute_multi_machine_logs( + self, + *, + follow: bool, + tail: Optional[int], + rank: Optional[int], + timestamps: bool, + since: Optional[str], + until: Optional[str], + query: Optional[str], + severity: Optional[str], + ) -> Union[str, Iterator[str]]: + if rank is not None: + raise ValueError( + "`rank` is not supported on a multi-machine parent; " + "read a single machine with `job.machines[rank].logs` " + "(or `mmt.machines[rank].logs` through the compatibility API)." + ) + + status = self.status + if status not in (Status.Running, Status.Failed, Status.Completed, Status.Stopped): + raise RuntimeError(f"Logs are not available while the job is {status}.") + + lines = self._stream_multi_machine_entries( + follow=follow and status == Status.Running, + tail=tail, + timestamps=timestamps, + since=since, + until=until, + query=query, + severity=severity, + ) + if follow and status == Status.Running: + return lines + collected = list(lines) + return iter(collected) if follow else "\n".join(collected) + + def _stream_multi_machine_entries( + self, + *, + follow: bool, + tail: Optional[int], + timestamps: bool, + since: Optional[str], + until: Optional[str], + query: Optional[str], + severity: Optional[str], + ) -> Iterator[str]: + names = {machine._guaranteed_job.id: machine.name for machine in self.machines} + entries = self._logs_api.stream( + self.teamspace.id, + mmt_id=self._guaranteed_job.id, + since=since, + until=until, + query=query, + severity=severity, + follow=follow, + tail=tail, + tail_anchor=getattr(self._guaranteed_job, "stopped_at", None), + idle_timeout=None if follow else _RUNNING_LOGS_IDLE_TIMEOUT, + fallback_to_live=not follow, + stop=lambda: self.status in (Status.Stopped, Status.Completed, Status.Failed), + ) + for entry in entries: + yield entry.format(timestamps=timestamps, prefix=names.get(entry.resource_id, entry.resource_id)) + def _compute_logs_ranked( self, *, @@ -706,6 +861,9 @@ def _stream_logs( @property def link(self) -> str: + if self.is_multi_machine: + return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=mmt" + mmt_name = self._job_api.get_mmt_name(self._guaranteed_job) if self._job_api.get_image_name(self._guaranteed_job): @@ -728,6 +886,11 @@ def link(self) -> str: def image(self) -> Optional[str]: return self._job_api.get_image_name(self._guaranteed_job) + @property + def studio_name(self) -> Optional[str]: + """The name of the studio this job runs in, without instantiating the Studio.""" + return self._job_api.get_studio_name(self._guaranteed_job) + @property def studio(self) -> Optional["Studio"]: from lightning_sdk.studio import Studio @@ -743,10 +906,25 @@ def command(self) -> str: def _update_internal_job(self) -> None: if getattr(self, "_job", None) is None: - self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + if self.is_multi_machine: + self._attach_job(self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id)) + return + + from lightning_sdk.lightning_cloud.openapi.rest import ApiException + + try: + self._attach_job( + self._standalone_job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) + ) + except ApiException as ex: + if ex.status != 404: + raise + # Switch to the multi-machine API, then sync the real machine count from the payload. + self._num_machines = 2 + self._attach_job(self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id)) return - self._job = self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id) + self._attach_job(self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id)) @property def name(self) -> str: diff --git a/python/lightning_sdk/mmt.py b/python/lightning_sdk/mmt.py index 8b3fede0..c92d02ff 100644 --- a/python/lightning_sdk/mmt.py +++ b/python/lightning_sdk/mmt.py @@ -1,120 +1,66 @@ -import warnings -from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Protocol, Tuple, TypedDict, Union +from typing import TYPE_CHECKING, Dict, Optional, Protocol, Union -from lightning_sdk.api.cloud_account_api import CloudAccountApi -from lightning_sdk.api.logs_api import LogsApi -from lightning_sdk.api.mmt_api import MMTApiV2 -from lightning_sdk.api.utils import AccessibleResource, _get_cloud_url, raise_access_error_if_not_allowed -from lightning_sdk.job import _RUNNING_LOGS_IDLE_TIMEOUT, _Logs +from lightning_sdk.job import Job, JobDict from lightning_sdk.status import Status -from lightning_sdk.utils.logging import TrackCallsMeta -from lightning_sdk.utils.resolve import ( - _get_org_id, - _resolve_default_cloud_account, - _resolve_teamspace, - _setup_logger, - in_studio, - skip_studio_setup, -) if TYPE_CHECKING: - from lightning_sdk.job import Job from lightning_sdk.machine import CloudProvider, Machine from lightning_sdk.organization import Organization from lightning_sdk.studio import Studio from lightning_sdk.teamspace import Teamspace from lightning_sdk.user import User -_logger = _setup_logger(__name__) - __all__ = ["MMT", "MMTMachine"] -class MachineDict(TypedDict): - name: str - status: Status - machine: Union["Machine", str] - - class MMTMachine(Protocol): - """A single machine in multi-machine training.""" + """A single machine in a multi-machine job.""" @property def name(self) -> str: - """The name of the individual machine. Usually corresponds to the rank. - - Returns: - str: The name of this machine instance. - """ ... @property def machine(self) -> Union["Machine", str]: - """The actual machine type this node is running on. - - Returns: - Union[Machine, str]: The machine type of this node. - """ ... @property def artifact_path(self) -> Optional[str]: - """The path to the artifacts of this job. - - Returns: - Optional[str]: The artifact path, or None if not available. - """ ... @property def status(self) -> Status: - """The status of this job. - - Returns: - Status: The current status of this machine's job. - """ ... @property def resource_id(self) -> Optional[str]: - """The stable resource identifier for this machine.""" ... @property def private_ip_address(self) -> Optional[str]: - """The private IP address for this machine, if assigned.""" ... @property def placement_group_id(self) -> Optional[str]: - """The placement group identifier for this machine, if assigned.""" ... @property def rank(self) -> Optional[int]: - """The stable rank for this machine inside the multi-machine job.""" ... @property def logs(self) -> str: - """The logs of the given machine. - - Returns: - str: The complete logs from this machine's execution. - """ ... - def dict(self) -> MachineDict: - """Dict representation of the given machine. - - Returns: - MachineDict: A dictionary containing the machine's name, status, and machine type. - """ + def dict(self) -> JobDict: ... -class MMT(metaclass=TrackCallsMeta): - """Submit and manage multi-machine jobs on the Lightning AI Platform.""" +class MMT(Job): + """Compatibility interface for multi-machine jobs. + + Multi-machine functionality is implemented by :class:`lightning_sdk.job.Job`. + """ def __init__( self, @@ -124,42 +70,26 @@ def __init__( user: Union[str, "User", None] = None, *, _fetch_job: bool = True, + _num_machines: int = 2, ) -> None: - """Fetch already existing multi-machine jobs. - - Args: - name: the name of the job. - teamspace: the teamspace the job is part of. - org: the name of the organization owning the ``teamspace`` in case it is owned by an org. - Deprecated — pass the owner as part of ``teamspace`` instead, e.g. ``teamspace="owner/teamspace"``. - user: the name of the user owning the ``teamspace`` in case it is owned directly by a user instead - of an org. Deprecated — pass the owner as part of ``teamspace`` instead, - e.g. ``teamspace="owner/teamspace"``. - - Raises: - ValueError: If the teamspace cannot be resolved from the provided arguments, or if the job is not found - when ``_fetch_job=True``. - PermissionError: If the user does not have access to jobs in the given teamspace. - """ - teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user) - if teamspace is None: - raise ValueError( - "Cannot resolve the teamspace from provided arguments." - f" Got teamspace={teamspace}, org={org}, user={user}." + try: + super().__init__( + name=name, + teamspace=teamspace, + org=org, + user=user, + _fetch_job=_fetch_job, + # Default 2 forces the multi-machine API for lookup; real count is synced after fetch/attach. + _num_machines=_num_machines, ) - - raise_access_error_if_not_allowed(AccessibleResource.Jobs, teamspace_id=teamspace.id) - - self._teamspace = teamspace - self._name = name - self._job = None - self._prevent_refetch_latest = False - self._cloud_account_api = CloudAccountApi() - self._job_api = MMTApiV2() - self._logs_api = LogsApi() - - if _fetch_job: - self._update_internal_job() + except ValueError as ex: + # Job.__init__ raises "Job {name} does not exist…" on 404; keep the MMT-specific + # wording for that case only. Propagate teamspace/validation errors unchanged. + if "does not exist in Teamspace" not in str(ex): + raise + resolved_teamspace = getattr(self, "_teamspace", None) + teamspace_name = getattr(resolved_teamspace, "name", teamspace) + raise ValueError(f"Multi-machine job {name} does not exist in Teamspace {teamspace_name}") from ex @classmethod def run( @@ -184,135 +114,21 @@ def run( reuse_snapshot: bool = True, placement_group_id: Optional[str] = None, ) -> "MMT": - """Run async workloads using a docker image across multiple machines. - - Args: - name: The name of the job. Needs to be unique within the teamspace. - num_machines: The number of machines to run on. - machine: The machine type to run the job on. - command: The command to run inside your job. Required if using a studio. Optional if using an image. - If not provided for images, will run the container entrypoint and default command. - studio: The studio env to run the job with. Mutually exclusive with image. - image: The docker image to run the job with. Mutually exclusive with studio. - teamspace: The teamspace the job should be associated with. Defaults to the current teamspace. - Accepts a bare name or an ``owner/teamspace`` slug. - org: The organization owning the teamspace, if any. Defaults to the current organization. - Deprecated — pass the owner as part of ``teamspace`` instead, e.g. ``teamspace="owner/teamspace"``. - user: The user owning the teamspace, if any. Defaults to the current user. - Deprecated — pass the owner as part of ``teamspace`` instead, e.g. ``teamspace="owner/teamspace"``. - cloud: Cloud provider or cloud account to run the job on. - env: Environment variables to set inside the job. - interruptible: Whether the job should run on interruptible instances. Cheaper but can be preempted. - image_credentials: Credentials secret name used to pull a private image. - cloud_account_auth: Whether to authenticate with the cloud account to pull the image. - Required if the registry is part of a cloud provider, such as ECR. - entrypoint: The entrypoint of your docker container. Defaults to ``sh -c`` which - just runs the provided command in a standard shell if a command is provided. - If no command is provided, it will run the pre-defined entrypoint of the provided image. - To use the pre-defined entrypoint of the provided image with a specified command, - set this to an empty string. - Only applicable when submitting docker jobs. - path_mappings: Maps container paths to data-connection paths in the form - ``{"": ":"}`` or ``{"": ""}`` - for the root of a connection. Only applicable when submitting docker jobs. - max_runtime: Duration in seconds to allocate the machine. Required for some top-end GCP machines. - Defaults to 3 hours. - reuse_snapshot: Whether to reuse a Studio snapshot when multiple jobs for the same Studio are - submitted. Turning this off may result in longer startup times. Defaults to True. - placement_group_id: Optional placement group identifier for colocating the job. - - Returns: - MMT: The newly submitted multi-machine job instance. - - Raises: - ValueError: If required arguments are missing or mutually exclusive arguments are both provided. - RuntimeError: If image and studio are both provided. - """ - from lightning_sdk.lightning_cloud.openapi.rest import ApiException - from lightning_sdk.studio import Studio - - cloud_account = _resolve_default_cloud_account(None) - if cloud is not None: - cloud_account = None - if num_machines <= 1: raise ValueError("Multi-Machine training cannot be run with less than 2 Machines") - if not name: - raise ValueError("A job needs to have a name!") - - if image is None: - if not isinstance(studio, Studio): - with skip_studio_setup(): - studio = Studio( - name=studio, - teamspace=teamspace, - org=org, - user=user, - cloud=cloud, - create_ok=False, - ) - - if teamspace is None: - teamspace = studio.teamspace - else: - teamspace_name = teamspace if isinstance(teamspace, str) else teamspace.name - if studio.teamspace.name != teamspace_name: - raise ValueError( - "Studio teamspace does not match provided teamspace. " - "Can only run jobs with Studio envs in the teamspace of that Studio." - ) - - if cloud_account is None: - cloud_account = studio.cloud_account - - if cloud_account != studio.cloud_account: - raise ValueError( - "Studio cloud_account does not match provided cloud_account. " - "Can only run jobs with Studio envs in the same cloud_account." - ) - - if image_credentials is not None: - raise ValueError("image_credentials is only supported when using a custom image") - - if cloud_account_auth: - raise ValueError("cloud_account_auth is only supported when using a custom image") - - if entrypoint is not None: - raise ValueError("Specifying the entrypoint has no effect for jobs with Studio envs.") - - else: - if studio is not None: - raise RuntimeError( - "image and studio are mutually exclusive as both define the environment to run the job in" - ) - - if cloud_account is None and cloud is None and in_studio(): - try: - with skip_studio_setup(): - resolve_studio = Studio(teamspace=teamspace, user=user, org=org) - cloud_account = resolve_studio.cloud_account - except (ValueError, ApiException): - warnings.warn("Could not infer cloud account from studio. Using teamspace default.") - - if command is not None and entrypoint is None: - entrypoint = "sh -c" - elif entrypoint == "" or entrypoint is None: - entrypoint = None - - mmt = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False) - submit_cloud = cloud if cloud_account is None else None - - mmt._submit( - num_machines=num_machines, + return super().run( + name=name, machine=machine, - cloud=submit_cloud, + cloud=cloud, command=command, studio=studio, image=image, + teamspace=teamspace, + org=org, + user=user, env=env, interruptible=interruptible, - cloud_account=cloud_account, image_credentials=image_credentials, cloud_account_auth=cloud_account_auth, entrypoint=entrypoint, @@ -320,348 +136,5 @@ def run( max_runtime=max_runtime, reuse_snapshot=reuse_snapshot, placement_group_id=placement_group_id, - ) - - _logger.info(f"Multi-Machine Job was successfully launched. View it at {mmt.link}") - return mmt - - def _submit( - self, - num_machines: int, - machine: Union["Machine", str], - cloud: Optional[Union["CloudProvider", str]] = None, - command: Optional[str] = None, - studio: Optional["Studio"] = None, - image: Optional[str] = None, - env: Optional[Dict[str, str]] = None, - interruptible: bool = False, - cloud_account: Optional[str] = None, - image_credentials: Optional[str] = None, - cloud_account_auth: bool = False, - entrypoint: Optional[str] = None, - path_mappings: Optional[Dict[str, str]] = None, - max_runtime: Optional[int] = None, - reuse_snapshot: bool = True, - placement_group_id: Optional[str] = None, - ) -> "MMT": - if studio is not None: - studio_id = studio._studio.id - if image is not None: - raise ValueError( - "image and studio are mutually exclusive as both define the environment to run the job in" - ) - if command is None: - raise ValueError("command is required when using a studio") - else: - studio_id = None - if image is None: - raise ValueError("either image or studio must be provided") - - cloud_account = self._cloud_account_api.resolve_cloud_account( - self._teamspace.id, - cloud=cloud or cloud_account, - default_cloud_account=self._teamspace.default_cloud_account, - ) - - submitted = self._job_api.submit_job( - name=self.name, num_machines=num_machines, - command=command, - cloud_account=cloud_account, - teamspace_id=self._teamspace.id, - studio_id=studio_id, - image=image, - machine=machine, - interruptible=interruptible, - env=env, - image_credentials=image_credentials, - cloud_account_auth=cloud_account_auth, - entrypoint=entrypoint, - path_mappings=path_mappings, - max_runtime=max_runtime, - reuse_snapshot=reuse_snapshot, - placement_group_id=placement_group_id, ) - self._job = submitted - self._name = submitted.name - return self - - @property - def machines(self) -> Tuple["Job", ...]: - from lightning_sdk.job import Job - - subjobs = sorted( - self._job_api.list_mmt_subjobs(self._guaranteed_job.id, self.teamspace.id), - key=lambda job: job.spec.rank, - ) - machines = [] - for subjob in subjobs: - job = Job(name=subjob.name, teamspace=self.teamspace, _fetch_job=False) - job._job = subjob - machines.append(job) - return tuple(machines) - - def stop(self) -> None: - if self.status in (Status.Stopped, Status.Completed, Status.Failed): - return - self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id) - - def delete(self) -> None: - self._job_api.delete_job( - job_id=self._guaranteed_job.id, - teamspace_id=self._teamspace.id, - ) - - def wait(self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False) -> None: - import time - - start = time.time() - while True: - if self.status in (Status.Completed, Status.Stopped, Status.Failed): - break - - if timeout is not None and time.time() - start > timeout: - if stop_on_timeout: - self.stop() - raise TimeoutError("Job didn't finish within the provided timeout.") - - time.sleep(interval) - - async def async_wait( - self, interval: float = 5.0, timeout: Optional[float] = None, stop_on_timeout: bool = False - ) -> None: - import asyncio - - start = asyncio.get_event_loop().time() - while True: - if self.status in (Status.Completed, Status.Stopped, Status.Failed): - break - - if timeout is not None and asyncio.get_event_loop().time() - start > timeout: - if stop_on_timeout: - self.stop() - raise TimeoutError("Job didn't finish within the provided timeout.") - - await asyncio.sleep(interval) - - @property - def status(self) -> Status: - return self._job_api._job_state_to_external(self._latest_job.state) - - @property - def id(self) -> Optional[str]: - """The multi-machine job's unique identifier.""" - return self._job.id if self._job is not None else None - - @property - def placement_group_id(self) -> Optional[str]: - return self._guaranteed_job.spec.placement_group_id - - @property - def artifact_path(self) -> Optional[str]: - raise NotImplementedError - - @property - def snapshot_path(self) -> Optional[str]: - raise NotImplementedError - - @property - def share_path(self) -> Optional[str]: - return None - - @property - def machine(self) -> Union["Machine", str]: - return self._job_api._get_job_machine_from_spec( - self._guaranteed_job.spec, - self.teamspace.id, - _get_org_id(self.teamspace), - ) - - def _update_internal_job(self) -> None: - if getattr(self, "_job", None) is None: - from lightning_sdk.lightning_cloud.openapi.rest import ApiException - - try: - self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id) - except ApiException as ex: - if ex.status != 404: - raise - raise ValueError( - f"Multi-machine job {self._name} does not exist in Teamspace {self._teamspace.name}" - ) from ex - return - - self._job = self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id) - - @property - def name(self) -> str: - return self._name - - @property - def resource_id(self) -> Optional[str]: - return self._guaranteed_job.id - - @property - def teamspace(self) -> "Teamspace": - return self._teamspace - - @property - def link(self) -> str: - return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=mmt" - - @property - def image(self) -> Optional[str]: - return self._job_api.get_image_name(self._guaranteed_job) - - @property - def studio(self) -> Optional["Studio"]: - from lightning_sdk.studio import Studio - - studio_name = self._job_api.get_studio_name(self._guaranteed_job) - if not studio_name: - return None - return Studio(studio_name, teamspace=self.teamspace) - - @property - def command(self) -> str: - return self._job_api.get_command(self._guaranteed_job) - - @property - def num_machines(self) -> int: - return self._job_api.get_num_machines(self._guaranteed_job) - - @property - def logs(self) -> _Logs: - """The logs of every machine, merged into one timeline. - - Use it as a value for a snapshot of the logs up to now:: - - print(mmt.logs) - - or call it to pass options and/or follow the logs live:: - - recent = mmt.logs(tail=100) # snapshot of the last 100 lines - for line in mmt.logs(follow=True): # stream new lines as they arrive - print(line) - - Options: - - - ``follow``: Keep the stream open and yield new lines as they are produced. - Returns an iterator of lines instead of a string. - - ``tail``: Only include the last N lines. - - ``timestamps``: Prepend each line with its ISO-8601 timestamp. - - ``since``/``until``: Only include lines within this RFC3339 time range. - - ``query``: Only include lines containing every whitespace-separated term. - - ``severity``: Only include lines at or above this level (``error``, ``warning``, - ``info`` or ``debug``). - - Every line is labelled with the machine it came from. To read a single machine, use - ``mmt.machines[rank].logs``. - """ - return _Logs(self._compute_logs) - - def _compute_logs( - self, - *, - follow: bool = False, - tail: Optional[int] = None, - rank: Optional[int] = None, - timestamps: bool = False, - since: Optional[str] = None, - until: Optional[str] = None, - query: Optional[str] = None, - severity: Optional[str] = None, - ) -> Union[str, Iterator[str]]: - """Fetch the merged logs of every machine. See :attr:`logs` for the public API.""" - if rank is not None: - raise ValueError("`rank` is not supported here; read a single machine with `mmt.machines[rank].logs`.") - - status = self.status - if status not in (Status.Running, Status.Failed, Status.Completed, Status.Stopped): - raise RuntimeError(f"Logs are not available while the job is {status}.") - - lines = self._stream_entries( - follow=follow and status == Status.Running, - tail=tail, - timestamps=timestamps, - since=since, - until=until, - query=query, - severity=severity, - ) - if follow and status == Status.Running: - return lines - collected = list(lines) - return iter(collected) if follow else "\n".join(collected) - - def _stream_entries( - self, - *, - follow: bool, - tail: Optional[int], - timestamps: bool, - since: Optional[str] = None, - until: Optional[str] = None, - query: Optional[str] = None, - severity: Optional[str] = None, - ) -> Iterator[str]: - """Yield formatted log lines for every machine, labelled with the machine they came from.""" - names = {machine._guaranteed_job.id: machine.name for machine in self.machines} - entries = self._logs_api.stream( - self.teamspace.id, - mmt_id=self._guaranteed_job.id, - since=since, - until=until, - query=query, - severity=severity, - follow=follow, - tail=tail, - # A finished job's last lines sit at its stop time, so start the tail search there - # instead of walking back from now through a job that ran days ago. - tail_anchor=getattr(self._guaranteed_job, "stopped_at", None), - idle_timeout=None if follow else _RUNNING_LOGS_IDLE_TIMEOUT, - # A running job whose logs are not in the current storage format yet has no saved - # history; tail its live stream so a snapshot still shows something. - fallback_to_live=not follow, - stop=lambda: self.status in (Status.Stopped, Status.Completed, Status.Failed), - ) - for entry in entries: - yield entry.format(timestamps=timestamps, prefix=names.get(entry.resource_id, entry.resource_id)) - - def dict(self) -> Dict[str, object]: - studio = self.studio - - return { - "name": self.name, - "teamspace": f"{self.teamspace.owner.name}/{self.teamspace.name}", - "studio": studio.name if studio else None, - "image": self.image, - "command": self.command, - "status": self.status, - "machine": self.machine, - "total_cost": self.total_cost, - } - - def json(self) -> str: - import json - - return json.dumps(self.dict(), indent=4, sort_keys=True, default=str) - - @property - def _guaranteed_job(self) -> Any: - if getattr(self, "_job", None) is None: - self._update_internal_job() - - return self._job - - @property - def total_cost(self) -> float: - return self._job_api.get_total_cost(self._latest_job) - - @property - def _latest_job(self) -> Any: - if self._prevent_refetch_latest: - return self._guaranteed_job - - self._update_internal_job() - return self._job diff --git a/python/lightning_sdk/pipeline/steps.py b/python/lightning_sdk/pipeline/steps.py index e41d3427..17fbf60f 100644 --- a/python/lightning_sdk/pipeline/steps.py +++ b/python/lightning_sdk/pipeline/steps.py @@ -207,6 +207,7 @@ def __init__( reuse_snapshot: bool = True, scratch_disks: Optional[Dict[str, int]] = None, placement_group_id: Optional[str] = None, + num_machines: int = 1, ) -> None: """Configure a job step in a pipeline. @@ -231,8 +232,11 @@ def __init__( reuse_snapshot: Whether to reuse a studio snapshot across jobs. Defaults to True. scratch_disks: Extra volumes to mount under ``/teamspace/scratch``. placement_group_id: Optional placement group identifier for colocating the job. + num_machines: Number of machines to allocate. Defaults to one. """ + if num_machines < 1: + raise ValueError("A job needs to run on at least one machine") self.name = name self.machine = machine or Machine.CPU self.command = command @@ -254,6 +258,7 @@ def __init__( self.reuse_snapshot = reuse_snapshot self.scratch_disks = scratch_disks self.placement_group_id = placement_group_id + self.num_machines = num_machines def to_proto( self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem] @@ -286,26 +291,36 @@ def to_proto( _validate_cloud_account(cloud_account, resolved_cloud_account, shared_filesystem) - body = JobApiV2._create_job_body( - name=self.name, - command=self.command, - cloud_account=resolved_cloud_account or cloud_account, - studio_id=studio._studio.id if isinstance(studio, Studio) else None, - image=self.image, - machine=self.machine, - interruptible=self.interruptible, - env=self.env, - image_credentials=self.image_credentials, - cloud_account_auth=self.cloud_account_auth, - entrypoint=self.entrypoint, - path_mappings=self.path_mappings, - max_runtime=self.max_runtime, - machine_image_version=machine_image_version, - reuse_snapshot=self.reuse_snapshot, - scratch_disks=self.scratch_disks, - placement_group_id=self.placement_group_id, - ) + body_kwargs = { + "name": self.name, + "command": self.command, + "cloud_account": resolved_cloud_account or cloud_account, + "studio_id": studio._studio.id if isinstance(studio, Studio) else None, + "image": self.image, + "machine": self.machine, + "interruptible": self.interruptible, + "env": self.env, + "image_credentials": self.image_credentials, + "cloud_account_auth": self.cloud_account_auth, + "entrypoint": self.entrypoint, + "path_mappings": self.path_mappings, + "max_runtime": self.max_runtime, + "machine_image_version": machine_image_version, + "reuse_snapshot": self.reuse_snapshot, + "placement_group_id": self.placement_group_id, + } + if self.num_machines > 1: + if self.scratch_disks: + raise ValueError("scratch_disks are not supported for multi-machine jobs") + body = MMTApiV2._create_mmt_body(num_machines=self.num_machines, **body_kwargs) + return V1PipelineStep( + name=self.name, + type=V1PipelineStepType.MMT, + wait_for=_to_wait_for(self.wait_for), + mmt=body, + ) + body = JobApiV2._create_job_body(scratch_disks=self.scratch_disks, **body_kwargs) return V1PipelineStep( name=self.name, type=V1PipelineStepType.JOB, diff --git a/python/lightning_sdk/teamspace.py b/python/lightning_sdk/teamspace.py index bc0d3ba8..668c3391 100644 --- a/python/lightning_sdk/teamspace.py +++ b/python/lightning_sdk/teamspace.py @@ -278,7 +278,7 @@ def clusters(self) -> List[str]: @property def jobs(self) -> Tuple["Job", ...]: - """All single-machine jobs in this teamspace. + """All standalone and multi-machine jobs in this teamspace. Returns: tuple[Job, ...]: Every Job belonging to this teamspace. @@ -294,7 +294,11 @@ def jobs(self) -> Tuple["Job", ...]: for j2 in self._teamspace_api.list_jobs(teamspace_id=self.id): # _fetch_job = False to prevent refetching on init since we already got it job = Job(name=j2.name, teamspace=self, _fetch_job=False) - job._job = j2 + job._attach_job(j2) + jobs.append(job) + for m2 in self._teamspace_api.list_mmts(teamspace_id=self.id): + job = Job(name=m2.name, teamspace=self, _fetch_job=False) + job._attach_job(m2) jobs.append(job) return tuple(jobs) @@ -317,7 +321,7 @@ def multi_machine_jobs(self) -> Tuple["MMT", ...]: for m2 in self._teamspace_api.list_mmts(teamspace_id=self.id): # _fetch_job = False to prevent refetching on init since we already got it mmt = MMT(name=m2.name, teamspace=self, _fetch_job=False) - mmt._job = m2 + mmt._attach_job(m2) mmts.append(mmt) return tuple(mmts) diff --git a/python/tests/cli/job/test_inspect.py b/python/tests/cli/job/test_inspect.py index b8081967..559f4ea8 100644 --- a/python/tests/cli/job/test_inspect.py +++ b/python/tests/cli/job/test_inspect.py @@ -2,6 +2,7 @@ from click.testing import CliRunner +from lightning_sdk.mmt import MMT from tests.cli.help import assert_help_contains, mock_command_logging @@ -65,6 +66,29 @@ def test_job_inspect_uses_positional_name() -> None: job.json.assert_called_once_with() +@mock_command_logging +def test_job_inspect_selects_mmt_rank() -> None: + from lightning_sdk.cli.job.inspect import inspect_job + + teamspace = MagicMock() + mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True + machine = MagicMock() + machine.json.return_value = '{"name":"distributed-1"}' + with patch("lightning_sdk.cli.job.inspect.resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.cli.job.inspect.resolve_job", + return_value=mmt, + ), patch( + "lightning_sdk.cli.job.inspect.resolve_job_machine", + return_value=machine, + ) as resolve_rank: + result = CliRunner().invoke(inspect_job, ["distributed", "--rank", "1"]) + + assert result.exit_code == 0 + resolve_rank.assert_called_once_with(mmt, 1) + assert "distributed-1" in result.output + + @mock_command_logging def test_job_inspect_help_shows_positional_name() -> None: assert_help_contains("lightning job inspect --help", "Usage: lightning job inspect [OPTIONS] [NAME]") diff --git a/python/tests/cli/job/test_list.py b/python/tests/cli/job/test_list.py index d3df143e..fdd8b84a 100644 --- a/python/tests/cli/job/test_list.py +++ b/python/tests/cli/job/test_list.py @@ -1,5 +1,10 @@ +import json +from types import SimpleNamespace +from unittest.mock import patch + from click.testing import CliRunner +from lightning_sdk.cli.job.list import list_jobs from lightning_sdk.cli.legacy.list import jobs from tests.cli.help import assert_help_contains, mock_command_logging @@ -21,6 +26,60 @@ def test_jobs_list_help() -> None: assert_help_contains("lightning jobs list --help", "Usage: lightning jobs list", "List jobs for a given teamspace.") +def _teamspace_with_jobs() -> SimpleNamespace: + """Build a teamspace whose jobs expose only the attributes real Job/MMT objects have.""" + owner = SimpleNamespace(name="org") + teamspace = SimpleNamespace(name="teamspace", owner=owner) + single = SimpleNamespace( + name="single", + teamspace=teamspace, + studio_name=None, + image="ubuntu", + status="Running", + machine="CPU", + total_cost=1.0, + ) + multi = SimpleNamespace( + name="distributed", + teamspace=teamspace, + studio_name=None, + image="ubuntu", + status="Running", + machine="CPU", + num_machines=4, + total_cost=4.0, + ) + teamspace.jobs = [single, multi] + teamspace.multi_machine_jobs = [multi] + return teamspace + + +@mock_command_logging +def test_job_list_includes_single_and_multi_machine_jobs() -> None: + teamspace = _teamspace_with_jobs() + + with patch("lightning_sdk.cli.job.list.resolve_teamspace", return_value=teamspace): + result = CliRunner().invoke(list_jobs, ["--json"]) + + assert result.exit_code == 0, result.output + rows = json.loads(result.output) + assert [(row["name"], row["num_machines"]) for row in rows] == [ + ("distributed", 4), + ("single", 1), + ] + + +@mock_command_logging +def test_job_list_sort_by_cloud_account_without_attribute() -> None: + teamspace = _teamspace_with_jobs() + + with patch("lightning_sdk.cli.job.list.resolve_teamspace", return_value=teamspace): + result = CliRunner().invoke(list_jobs, ["--sort-by", "cloud-account", "--json"]) + + assert result.exit_code == 0, result.output + assert {row["name"] for row in json.loads(result.output)} == {"single", "distributed"} + + @mock_command_logging def test_list_jobs_legacy_help() -> None: assert_help_contains( diff --git a/python/tests/cli/job/test_logs.py b/python/tests/cli/job/test_logs.py index 59c34ec3..34578441 100644 --- a/python/tests/cli/job/test_logs.py +++ b/python/tests/cli/job/test_logs.py @@ -2,6 +2,7 @@ from click.testing import CliRunner +from lightning_sdk.mmt import MMT from tests.cli.help import assert_help_contains, mock_command_logging @@ -79,6 +80,52 @@ def test_job_logs_follows_with_options() -> None: ) +@mock_command_logging +def test_job_logs_selects_mmt_rank() -> None: + from lightning_sdk.cli.job.logs import logs_job + + mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True + machine = MagicMock() + machine.logs.return_value = "rank one" + with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.logs.resolve_job", + return_value=mmt, + ), patch( + "lightning_sdk.cli.job.logs.resolve_job_machine", + return_value=machine, + ) as resolve_rank: + result = CliRunner().invoke(logs_job, ["distributed", "--rank", "1"]) + + assert result.exit_code == 0 + assert "rank one" in result.output + resolve_rank.assert_called_once_with(mmt, 1) + machine.logs.assert_called_once_with( + follow=False, tail=None, rank=0, timestamps=False, since=None, until=None, query=None, severity=None + ) + + +@mock_command_logging +def test_job_logs_merges_mmt_without_rank() -> None: + from lightning_sdk.cli.job.logs import logs_job + + mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True + mmt.logs.return_value = "[distributed-0] zero\n[distributed-1] one" + with patch("lightning_sdk.cli.job.logs.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.logs.resolve_job", + return_value=mmt, + ): + result = CliRunner().invoke(logs_job, ["distributed"]) + + assert result.exit_code == 0 + assert "distributed-0" in result.output + assert "distributed-1" in result.output + mmt.logs.assert_called_once_with( + follow=False, tail=None, timestamps=False, since=None, until=None, query=None, severity=None + ) + + @mock_command_logging def test_job_logs_passes_filters() -> None: from lightning_sdk.cli.job.logs import logs_job diff --git a/python/tests/cli/job/test_run.py b/python/tests/cli/job/test_run.py index c1078a08..243b3c48 100644 --- a/python/tests/cli/job/test_run.py +++ b/python/tests/cli/job/test_run.py @@ -69,6 +69,24 @@ def test_run_job_with_cloud(monkeypatch): assert mock_job.run.call_args.kwargs["cloud"] == "aws" +@mock_command_logging +def test_run_job_with_multiple_machines_uses_job() -> None: + submitted = MagicMock(name="distributed") + submitted.name = "distributed" + with patch("lightning_sdk.cli.job.run.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.run.Job.run", + return_value=submitted, + ) as run_sdk_job: + result = CliRunner().invoke( + run_job, + ["--name", "distributed", "--image", "ubuntu", "--num-machines", "3"], + ) + + assert result.exit_code == 0, result.output + assert run_sdk_job.call_args.kwargs["num_machines"] == 3 + assert "Submitted job distributed." in result.output + + @pytest.mark.parametrize( ("input_mappings", "expected"), [ diff --git a/python/tests/cli/job/test_ssh.py b/python/tests/cli/job/test_ssh.py index 53610eb3..cc09d637 100644 --- a/python/tests/cli/job/test_ssh.py +++ b/python/tests/cli/job/test_ssh.py @@ -5,6 +5,7 @@ from click.testing import CliRunner from lightning_sdk.cli.job.ssh import _ssh_user_for_job_id, ssh_impl, ssh_job +from lightning_sdk.mmt import MMT from lightning_sdk.status import Status from tests.cli.help import assert_help_contains, command_text, mock_command_logging @@ -27,7 +28,7 @@ def test_job_ssh_help() -> None: assert "Usage: lightning job ssh [OPTIONS] NAME" in result_text assert "SSH into a running job." in result_text assert "--teamspace" in result_text - assert "--rank" not in result_text + assert "--rank" in result_text @mock_command_logging @@ -88,6 +89,30 @@ def test_ssh_runs_against_job_gateway_user() -> None: run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_01jj4hvvjj4zx1t1esm5az3zt7@ssh.lightning.ai"]) +def test_ssh_selects_mmt_rank() -> None: + mmt = MagicMock(spec=MMT) + mmt.is_multi_machine = True + machine = MagicMock() + machine.name = "distributed-1" + machine.status = Status.Running + machine.id = "job_rank1" + + with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.job.ssh.resolve_job", + return_value=mmt, + ), patch( + "lightning_sdk.cli.job.ssh.resolve_job_machine", + return_value=machine, + ) as resolve_rank, patch( + "lightning_sdk.cli.job.ssh.configure_ssh_internal", + return_value="/tmp/lightning_rsa", + ), patch("lightning_sdk.cli.job.ssh.subprocess.run") as run: + ssh_impl(name="distributed", teamspace=None, rank=1) + + resolve_rank.assert_called_once_with(mmt, 1) + run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_rank1@ssh.lightning.ai"]) + + def test_ssh_retries_with_fresh_keys_on_failure() -> None: job = MagicMock() job.name = "train" diff --git a/python/tests/cli/mmt/test_ssh.py b/python/tests/cli/mmt/test_ssh.py index cc3697c6..a15105ef 100644 --- a/python/tests/cli/mmt/test_ssh.py +++ b/python/tests/cli/mmt/test_ssh.py @@ -60,6 +60,7 @@ def test_ssh_resolves_before_downloading_keys() -> None: def test_ssh_rejects_non_running_machine() -> None: rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 rank0.status = Status.Completed rank0.id = "job_01abc" @@ -81,11 +82,13 @@ def test_ssh_defaults_to_rank_zero() -> None: teamspace = MagicMock() rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 rank0.status = Status.Running rank0.id = "job_rank0" rank1 = MagicMock() rank1.name = "train-1" + rank1.rank = 1 rank1.status = Status.Running rank1.id = "job_rank1" @@ -109,11 +112,13 @@ def test_ssh_selects_requested_rank() -> None: teamspace = MagicMock() rank0 = MagicMock() rank0.name = "olmo3-7b-think-sft-full-0" + rank0.rank = 0 rank0.status = Status.Running rank0.id = "job_r0" rank1 = MagicMock() rank1.name = "olmo3-7b-think-sft-full-1" + rank1.rank = 1 rank1.status = Status.Running rank1.id = "job_r1" @@ -135,6 +140,7 @@ def test_ssh_selects_requested_rank() -> None: def test_ssh_rejects_unknown_rank() -> None: rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 mmt = MagicMock() mmt.name = "train" @@ -142,13 +148,42 @@ def test_ssh_rejects_unknown_rank() -> None: with patch("lightning_sdk.cli.mmt.ssh.resolve_teamspace", return_value=MagicMock()), patch( "lightning_sdk.cli.mmt.ssh.resolve_mmt", return_value=mmt - ), pytest.raises(click.ClickException, match="Rank 3 not found"): + ), pytest.raises(click.ClickException, match="Rank 3 not found.*Available ranks: 0"): ssh_impl(name="train", teamspace=None, rank=3) +def test_ssh_selects_rank_when_names_lack_rank_suffix() -> None: + """Machines are matched on their rank, not on a ``{job}-{rank}`` name convention.""" + rank0 = MagicMock() + rank0.name = "worker-alpha" + rank0.rank = 0 + rank0.status = Status.Running + rank0.id = "job_r0" + + rank1 = MagicMock() + rank1.name = "worker-beta" + rank1.rank = 1 + rank1.status = Status.Running + rank1.id = "job_r1" + + mmt = MagicMock() + mmt.name = "train" + mmt.machines = (rank0, rank1) + + with patch("lightning_sdk.cli.mmt.ssh.resolve_teamspace", return_value=MagicMock()), patch( + "lightning_sdk.cli.mmt.ssh.resolve_mmt", return_value=mmt + ), patch("lightning_sdk.cli.mmt.ssh.configure_ssh_internal", return_value="/tmp/lightning_rsa"), patch( + "lightning_sdk.cli.mmt.ssh.subprocess.run" + ) as run: + ssh_impl(name="train", teamspace=None, rank=1) + + run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_r1@ssh.lightning.ai"]) + + def test_ssh_retries_with_fresh_keys_on_failure() -> None: rank0 = MagicMock() rank0.name = "train-0" + rank0.rank = 0 rank0.status = Status.Running rank0.id = "job_01abc" diff --git a/python/tests/cli/utils/test_resource_resolution.py b/python/tests/cli/utils/test_resource_resolution.py index 029faa86..2c024000 100644 --- a/python/tests/cli/utils/test_resource_resolution.py +++ b/python/tests/cli/utils/test_resource_resolution.py @@ -9,6 +9,7 @@ resolve_cluster, resolve_deployment, resolve_job, + resolve_job_machine, resolve_mmt, resolve_studio, resolve_teamspace, @@ -184,6 +185,32 @@ def test_resolve_job_converts_not_found_to_usage_error() -> None: resolve_job("train", MagicMock()) +def test_resolve_job_machine_prefers_rank_attribute() -> None: + rank0 = MagicMock(name="odd-name-0") + rank0.name = "odd-name-0" + rank0.rank = 0 + rank1 = MagicMock(name="odd-name-1") + rank1.name = "odd-name-1" + rank1.rank = 1 + mmt = MagicMock(name="distributed", machines=(rank0, rank1)) + mmt.name = "distributed" + + assert resolve_job_machine(mmt, 1) is rank1 + + +def test_resolve_job_machine_falls_back_to_name_suffix() -> None: + rank0 = MagicMock(name="rank-0") + rank0.name = "distributed-0" + rank0.rank = None + rank1 = MagicMock(name="rank-1") + rank1.name = "distributed-1" + rank1.rank = None + mmt = MagicMock(name="distributed", machines=(rank0, rank1)) + mmt.name = "distributed" + + assert resolve_job_machine(mmt, 1) is rank1 + + def test_resolve_mmt_requires_name() -> None: with pytest.raises(click.UsageError, match="JOB"): resolve_mmt(None, MagicMock()) @@ -211,10 +238,10 @@ def test_resolve_mmt_converts_not_found_to_usage_error() -> None: def test_resolve_mmt_converts_real_lookup_not_found_to_usage_error() -> None: teamspace = MagicMock() teamspace.id = "teamspace-id" - with patch("lightning_sdk.mmt._resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.mmt.raise_access_error_if_not_allowed" - ), patch("lightning_sdk.mmt.CloudAccountApi"), patch( - "lightning_sdk.mmt.MMTApiV2.get_job_by_name", + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.raise_access_error_if_not_allowed" + ), patch("lightning_sdk.job.CloudAccountApi"), patch( + "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404, reason="Not Found"), ), pytest.raises(click.UsageError, match="distributed"): resolve_mmt("distributed", teamspace) @@ -222,7 +249,7 @@ def test_resolve_mmt_converts_real_lookup_not_found_to_usage_error() -> None: def test_resolve_mmt_preserves_pre_lookup_not_found_error() -> None: failure = ApiException(status=404, reason="Teamspace Not Found") - with patch("lightning_sdk.mmt._resolve_teamspace", side_effect=failure), pytest.raises(ApiException) as raised: + with patch("lightning_sdk.job._resolve_teamspace", side_effect=failure), pytest.raises(ApiException) as raised: resolve_mmt("distributed", MagicMock()) assert raised.value is failure @@ -232,10 +259,10 @@ def test_resolve_mmt_preserves_real_lookup_permission_error() -> None: teamspace = MagicMock() teamspace.id = "teamspace-id" failure = ApiException(status=403, reason="Forbidden") - with patch("lightning_sdk.mmt._resolve_teamspace", return_value=teamspace), patch( - "lightning_sdk.mmt.raise_access_error_if_not_allowed" - ), patch("lightning_sdk.mmt.CloudAccountApi"), patch( - "lightning_sdk.mmt.MMTApiV2.get_job_by_name", + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.raise_access_error_if_not_allowed" + ), patch("lightning_sdk.job.CloudAccountApi"), patch( + "lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=failure, ), pytest.raises(ApiException) as raised: resolve_mmt("distributed", teamspace) diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index 25018fdb..92abdb82 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -11,6 +11,7 @@ V1Job, V1JobSpec, ) +from lightning_sdk.lightning_cloud.openapi.rest import ApiException from lightning_sdk.machine import Machine from lightning_sdk.status import Status from lightning_sdk.studio import Studio @@ -42,7 +43,9 @@ def init_job_error(): studio = Studio("st-abc", "ts-abc", org="org-abc") Job("xyz", studio.teamspace) - with pytest.raises(ValueError, match="Job xyz does not exist in Teamspace ts-abc"): + with mock.patch("lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404)), pytest.raises( + ValueError, match="Job xyz does not exist in Teamspace ts-abc" + ): init_job_error() @@ -109,7 +112,9 @@ def test_delete_job( with pytest.raises(RuntimeError, match="Job j-abc does not exist in Teamspace ts-abc. Did you delete it?"): job.status # noqa: B018 - with pytest.raises(ValueError, match="Job j-abc does not exist in Teamspace ts-abc"): + with mock.patch("lightning_sdk.job.MMTApiV2.get_job_by_name", side_effect=ApiException(status=404)), pytest.raises( + ValueError, match="Job j-abc does not exist in Teamspace ts-abc" + ): Job("j-abc", studio.teamspace) @@ -154,6 +159,7 @@ def test_submit_job_v2_image(internal_studio_init_mocker, machine, command, env, reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) @@ -240,6 +246,7 @@ def test_submit_job_v2_studio(internal_studio_init_mocker, machine, env, interru reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) @@ -605,9 +612,10 @@ def test_submit_jobv2_studio_resolve( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock - - Job.run("test-job", machine=Machine.CPU, command="echo hello", studio="st-abc", teamspace="ts-abc", org="org-abc") + with mock.patch.object(Job, "_submit", submit_mock): + Job.run( + "test-job", machine=Machine.CPU, command="echo hello", studio="st-abc", teamspace="ts-abc", org="org-abc" + ) submit_mock.assert_called_once_with( command="echo hello", @@ -626,6 +634,7 @@ def test_submit_jobv2_studio_resolve( reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) @@ -685,11 +694,10 @@ def test_submit_job_v2_image_from_studio( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock keeping_alive_mock = mock.MagicMock() StudioApi.start_keeping_alive = keeping_alive_mock - with mock.patch.dict( + with mock.patch.object(Job, "_submit", submit_mock), mock.patch.dict( os.environ, { "LIGHTNING_CLOUD_SPACE_ID": "st-abc", @@ -724,6 +732,7 @@ def test_submit_job_v2_image_from_studio( reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) assert keeping_alive_mock.call_count == 0 @@ -740,17 +749,16 @@ def test_run_job_with_cloud_provider( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock - - Job.run( - "test-job", - machine=Machine.CPU, - command="echo hello", - image="nginx", - teamspace="ts-abc", - org="org-abc", - cloud="nebius", - ) + with mock.patch.object(Job, "_submit", submit_mock): + Job.run( + "test-job", + machine=Machine.CPU, + command="echo hello", + image="nginx", + teamspace="ts-abc", + org="org-abc", + cloud="nebius", + ) submit_mock.assert_called_once_with( command="echo hello", @@ -769,6 +777,7 @@ def test_run_job_with_cloud_provider( reuse_snapshot=True, scratch_disks=None, placement_group_id=None, + num_machines=1, ) @@ -944,11 +953,10 @@ def test_submit_job_from_running_studio( from lightning_sdk.job import Job submit_mock = mock.MagicMock() - Job._submit = submit_mock keeping_alive_mock = mock.MagicMock() StudioApi.start_keeping_alive = keeping_alive_mock - with mock.patch.dict( + with mock.patch.object(Job, "_submit", submit_mock), mock.patch.dict( os.environ, { "LIGHTNING_CLOUD_SPACE_ID": "st-abc", diff --git a/python/tests/core/test_mmt.py b/python/tests/core/test_mmt.py index 6d60d26d..c7ad87e2 100644 --- a/python/tests/core/test_mmt.py +++ b/python/tests/core/test_mmt.py @@ -57,6 +57,7 @@ def test_submit_mmt_v2_image(internal_studio_init_mocker, machine, command, env, max_runtime=None, reuse_snapshot=True, placement_group_id=None, + scratch_disks=None, ) @@ -172,6 +173,7 @@ def test_submit_mmt_v2_studio(internal_studio_init_mocker, machine, env, interru max_runtime=None, reuse_snapshot=True, placement_group_id=None, + scratch_disks=None, ) @@ -473,7 +475,7 @@ def test_mmtv2_delete(mmt_api_get_job_by_name_mocker, internal_studio_init_mocke job.delete() - delete_job_mock.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001") + delete_job_mock.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001", cloudspace_id=None) @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) diff --git a/python/tests/core/test_permissions_integration.py b/python/tests/core/test_permissions_integration.py index 55c18e29..a8dc96b3 100644 --- a/python/tests/core/test_permissions_integration.py +++ b/python/tests/core/test_permissions_integration.py @@ -244,7 +244,7 @@ def test_job_init_raises_error_when_jobs_disabled(mock_teamspace_api, mock_resol # MMT class permission tests -@mock.patch("lightning_sdk.mmt._resolve_teamspace") +@mock.patch("lightning_sdk.job._resolve_teamspace") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @pytest.mark.project_permission_test() def test_mmt_init_raises_error_when_jobs_disabled(mock_teamspace_api, mock_resolve_teamspace): diff --git a/python/tests/core/test_pipeline.py b/python/tests/core/test_pipeline.py index 82203d1a..ecfa1ec5 100644 --- a/python/tests/core/test_pipeline.py +++ b/python/tests/core/test_pipeline.py @@ -171,6 +171,15 @@ def test_job_step_threads_placement_group_id(): assert proto.job.spec.placement_group_id == "pg-1" +@patch("lightning_sdk.pipeline.steps.CloudAccountApi", new=MagicMock()) +def test_job_step_supports_multiple_machines(): + job = JobStep(name="job-0", machine=Machine.CPU, num_machines=3) + proto = job.to_proto(MagicMock(), "", False) + + assert proto.type == V1PipelineStepType.MMT + assert proto.mmt.machines == 3 + + @pytest.mark.parametrize("interruption_retries", [0, 3]) @patch.object(teamspace, "TeamspaceApi", new=MagicMock()) @patch.object(pipeline_module, "_get_cluster", new=MagicMock()) diff --git a/python/tests/core/test_teamspace.py b/python/tests/core/test_teamspace.py index 1744ae3b..80b7f6d6 100644 --- a/python/tests/core/test_teamspace.py +++ b/python/tests/core/test_teamspace.py @@ -554,30 +554,39 @@ def test_download_model_version( ) +@mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_jobs") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_list_jobs( list_jobs_mock, + list_mmts_mock, internal_get_org_api_mocker, internal_teamspace_api_mocker, internal_user_api_mocker, ): jobs = [V1Job(name="jobv2-1"), V1Job(name="jobv2-2"), V1Job(name="jobv2-3")] + mmts = [V1MultiMachineJob(name="mmtv2-1")] ts = Teamspace("ts-abc", org="org-abc") list_jobs_mock.return_value = jobs + list_mmts_mock.return_value = mmts # it's important that there are no additional calls to fetch individual jobs here. # they'd raise API Errors since we only mock the teamspace APIs listing # and not individual fetch requests listed_jobs = ts.jobs - assert len(listed_jobs) == 3 + assert len(listed_jobs) == 4 assert all(isinstance(j, Job) for j in listed_jobs) for lj, jj in zip(listed_jobs, jobs): assert lj.name == jj.name assert lj._job is jj + assert not lj.is_multi_machine + + assert listed_jobs[-1].name == "mmtv2-1" + assert listed_jobs[-1]._job is mmts[0] + assert listed_jobs[-1].is_multi_machine @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") diff --git a/python/tests/core/test_unified_job.py b/python/tests/core/test_unified_job.py new file mode 100644 index 00000000..6f13a13c --- /dev/null +++ b/python/tests/core/test_unified_job.py @@ -0,0 +1,147 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from lightning_sdk.job import Job +from lightning_sdk.lightning_cloud.openapi import V1Job, V1JobSpec, V1MultiMachineJob +from lightning_sdk.lightning_cloud.openapi.rest import ApiException +from lightning_sdk.mmt import MMT + + +def _teamspace() -> SimpleNamespace: + return SimpleNamespace( + id="teamspace-id", + name="teamspace", + owner=SimpleNamespace(name="owner"), + default_cloud_account="default-cloud", + ) + + +def test_job_lookup_falls_back_to_multi_machine() -> None: + teamspace = _teamspace() + standalone_api = MagicMock() + standalone_api.get_job_by_name.side_effect = ApiException(status=404) + multi_api = MagicMock() + multi_api.get_job_by_name.return_value = V1MultiMachineJob( + id="mmt-id", name="distributed", machines=2, spec=V1JobSpec() + ) + multi_api.get_num_machines.return_value = 2 + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2", return_value=standalone_api + ), patch("lightning_sdk.job.MMTApiV2", return_value=multi_api): + job = Job("distributed", teamspace) + + assert job.is_multi_machine + assert job.num_machines == 2 + standalone_api.get_job_by_name.assert_called_once() + multi_api.get_job_by_name.assert_called_once() + + +def test_job_lookup_prefers_standalone_on_name_collision() -> None: + teamspace = _teamspace() + standalone_api = MagicMock() + standalone_api.get_job_by_name.return_value = V1Job(id="job-id", name="train", spec=V1JobSpec()) + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2", return_value=standalone_api + ), patch("lightning_sdk.job.MMTApiV2") as multi_api: + job = Job("train", teamspace) + + assert not job.is_multi_machine + assert job.num_machines == 1 + assert job.machines == (job,) + multi_api.return_value.get_job_by_name.assert_not_called() + + +def test_job_run_routes_multi_machine_submission() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.submit_job.return_value = V1MultiMachineJob(id="mmt-id", name="distributed", machines=3, spec=V1JobSpec()) + multi_api.get_num_machines.return_value = 3 + cloud_api = MagicMock() + cloud_api.resolve_cloud_account.return_value = "cloud-id" + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job._resolve_default_cloud_account", return_value=None + ), patch("lightning_sdk.job.CloudAccountApi", return_value=cloud_api), patch( + "lightning_sdk.job.MMTApiV2", return_value=multi_api + ): + job = Job.run( + name="distributed", + machine="CPU", + cloud="aws", + image="ubuntu", + teamspace=teamspace, + num_machines=3, + ) + + assert type(job) is Job + assert job.is_multi_machine + assert job.num_machines == 3 + assert multi_api.submit_job.call_args.kwargs["num_machines"] == 3 + + +def test_mmt_run_returns_compatibility_subclass() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.submit_job.return_value = V1MultiMachineJob(id="mmt-id", name="distributed", machines=2, spec=V1JobSpec()) + cloud_api = MagicMock() + cloud_api.resolve_cloud_account.return_value = "cloud-id" + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job._resolve_default_cloud_account", return_value=None + ), patch("lightning_sdk.job.CloudAccountApi", return_value=cloud_api), patch( + "lightning_sdk.job.MMTApiV2", return_value=multi_api + ): + job = MMT.run( + name="distributed", + num_machines=2, + machine="CPU", + cloud="aws", + image="ubuntu", + teamspace=teamspace, + ) + + assert isinstance(job, MMT) + assert isinstance(job, Job) + assert job.is_multi_machine + + +def test_mmt_lookup_skips_standalone_api() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.get_job_by_name.return_value = V1MultiMachineJob( + id="mmt-id", name="distributed", machines=2, spec=V1JobSpec() + ) + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2" + ) as standalone_api, patch("lightning_sdk.job.MMTApiV2", return_value=multi_api): + job = MMT("distributed", teamspace) + + assert isinstance(job, MMT) + assert job.is_multi_machine + standalone_api.return_value.get_job_by_name.assert_not_called() + multi_api.get_job_by_name.assert_called_once() + + +def test_mmt_lookup_rewrites_only_not_found_errors() -> None: + teamspace = _teamspace() + multi_api = MagicMock() + multi_api.get_job_by_name.side_effect = ApiException(status=404) + + with patch("lightning_sdk.job._resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.job.JobApiV2" + ), patch("lightning_sdk.job.MMTApiV2", return_value=multi_api), pytest.raises( + ValueError, match="Multi-machine job distributed does not exist in Teamspace teamspace" + ): + MMT("distributed", teamspace) + + +def test_mmt_lookup_propagates_teamspace_resolution_errors() -> None: + with patch("lightning_sdk.job._resolve_teamspace", return_value=None), pytest.raises( + ValueError, match="Cannot resolve the teamspace" + ): + MMT("distributed", "owner/teamspace")