wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/sdk/launch/runner/{gcp_vertex.py → vertex_runner.py}

@@ -1,5 +1,5 @@
 import datetime
-import os
+import logging
 import shlex
 import time
 from typing import Any, Dict, Optional
@@ -10,17 +10,20 @@ if False:
 import yaml
 
 import wandb
-from wandb.errors import LaunchError
+from wandb.apis.internal import Api
 from wandb.util import get_module
 
 from .._project_spec import LaunchProject, get_entry_point_command
 from ..builder.abstract import AbstractBuilder
-from ..builder.build import construct_gcp_registry_uri, get_env_vars_dict
-from ..utils import LOG_PREFIX, PROJECT_SYNCHRONOUS, run_shell
+from ..builder.build import get_env_vars_dict
+from ..environment.gcp_environment import GcpEnvironment
+from ..utils import LOG_PREFIX, PROJECT_SYNCHRONOUS, LaunchError, run_shell
 from .abstract import AbstractRun, AbstractRunner, Status
 
 GCP_CONSOLE_URI = "https://console.cloud.google.com"
 
+_logger = logging.getLogger(__name__)
+
 
 class VertexSubmittedRun(AbstractRun):
     def __init__(self, job: Any) -> None:
@@ -57,12 +60,14 @@ class VertexSubmittedRun(AbstractRun):
 
     def get_status(self) -> Status:
         job_state = str(self._job.state)  # extract from type PipelineState
-        if job_state == "PipelineState.PIPELINE_STATE_SUCCEEDED":
+        if job_state == "JobState.JOB_STATE_SUCCEEDED":
             return Status("finished")
-        if job_state == "PipelineState.PIPELINE_STATE_FAILED":
+        if job_state == "JobState.JOB_STATE_FAILED":
             return Status("failed")
-        if job_state == "PipelineState.PIPELINE_STATE_RUNNING":
+        if job_state == "JobState.JOB_STATE_RUNNING":
             return Status("running")
+        if job_state == "JobState.JOB_STATE_PENDING":
+            return Status("starting")
         return Status("unknown")
 
     def cancel(self) -> None:
@@ -70,47 +75,37 @@ class VertexSubmittedRun(AbstractRun):
 
 
 class VertexRunner(AbstractRunner):
-    """Runner class, uses a project to create a VertexSubmittedRun"""
+    """Runner class, uses a project to create a VertexSubmittedRun."""
+
+    def __init__(
+        self, api: Api, backend_config: Dict[str, Any], environment: GcpEnvironment
+    ) -> None:
+        """Initialize a VertexRunner instance."""
+        super().__init__(api, backend_config)
+        self.environment = environment
 
     def run(
         self,
         launch_project: LaunchProject,
-        builder: AbstractBuilder,
-        registry_config: Dict[str, Any],
+        builder: Optional[AbstractBuilder],
     ) -> Optional[AbstractRun]:
-
+        """Run a Vertex job."""
         aiplatform = get_module(  # noqa: F811
             "google.cloud.aiplatform",
             "VertexRunner requires google.cloud.aiplatform to be installed",
         )
-
-        resource_args = launch_project.resource_args.get("gcp_vertex")
+        resource_args = launch_project.resource_args.get("vertex")
+        if not resource_args:
+            resource_args = launch_project.resource_args.get("gcp-vertex")
         if not resource_args:
             raise LaunchError(
                 "No Vertex resource args specified. Specify args via --resource-args with a JSON file or string under top-level key gcp_vertex"
             )
-        gcp_config = get_gcp_config(resource_args.get("gcp_config") or "default")
-        gcp_project = (
-            resource_args.get("gcp_project")
-            or gcp_config["properties"]["core"]["project"]
-        )
-        gcp_zone = resource_args.get("gcp_region") or gcp_config["properties"].get(
-            "compute", {}
-        ).get("zone")
-        gcp_region = "-".join(gcp_zone.split("-")[:2])
         gcp_staging_bucket = resource_args.get("staging_bucket")
         if not gcp_staging_bucket:
             raise LaunchError(
                 "Vertex requires a staging bucket for training and dependency packages in the same region as compute. Specify a bucket under key staging_bucket."
             )
-        gcp_artifact_repo = resource_args.get("artifact_repo")
-        if not gcp_artifact_repo:
-            raise LaunchError(
-                "Vertex requires an Artifact Registry repository for the Docker image. Specify a repo under key artifact_repo."
-            )
-        gcp_docker_host = (
-            resource_args.get("docker_host") or f"{gcp_region}-docker.pkg.dev"
-        )
         gcp_machine_type = resource_args.get("machine_type") or "n1-standard-4"
         gcp_accelerator_type = (
             resource_args.get("accelerator_type") or "ACCELERATOR_TYPE_UNSPECIFIED"
@@ -124,9 +119,10 @@ class VertexRunner(AbstractRunner):
         )
         service_account = resource_args.get("service_account")
         tensorboard = resource_args.get("tensorboard")
-
         aiplatform.init(
-            project=gcp_project, location=gcp_region, staging_bucket=gcp_staging_bucket
+            project=self.environment.project,
+            location=self.environment.region,
+            staging_bucket=gcp_staging_bucket,
         )
         synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
 
@@ -135,21 +131,13 @@ class VertexRunner(AbstractRunner):
         if launch_project.docker_image:
             image_uri = launch_project.docker_image
         else:
-
-            repository = construct_gcp_registry_uri(
-                gcp_artifact_repo,
-                gcp_project,
-                gcp_docker_host,
-            )
             assert entry_point is not None
+            assert builder is not None
             image_uri = builder.build_image(
                 launch_project,
-                repository,
                 entry_point,
             )
 
-        if not self.ack_run_queue_item(launch_project):
-            return None
         # TODO: how to handle this?
         entry_cmd = get_entry_point_command(entry_point, launch_project.override_args)
 
@@ -176,18 +164,19 @@ class VertexRunner(AbstractRunner):
             display_name=gcp_training_job_name, worker_pool_specs=worker_pool_specs
         )
 
-        submitted_run = VertexSubmittedRun(job)
-
-        # todo: support gcp dataset?
-
         wandb.termlog(
             f"{LOG_PREFIX}Running training job {gcp_training_job_name} on {gcp_machine_type}."
         )
 
-        # when sync is True, vertex blocks the main thread on job completion. when False, vertex returns a Future
-        # on this thread but continues to block the process on another thread. always set sync=False so we can get
-        # the job info (dependent on job._gca_resource)
-        job.run(service_account=service_account, tensorboard=tensorboard, sync=False)
+        if synchronous:
+            job.run(service_account=service_account, tensorboard=tensorboard, sync=True)
+        else:
+            job.submit(
+                service_account=service_account,
+                tensorboard=tensorboard,
+            )
+
+        submitted_run = VertexSubmittedRun(job)
 
         while not getattr(job._gca_resource, "name", None):
             # give time for the gcp job object to be created and named, this should only loop a couple times max
@@ -196,12 +185,6 @@ class VertexRunner(AbstractRunner):
         wandb.termlog(
             f"{LOG_PREFIX}View your job status and logs at {submitted_run.get_page_link()}."
         )
-
-        # hacky: if user doesn't want blocking behavior, kill both main thread and the background thread. job continues
-        # to run remotely. this obviously doesn't work if we need to do some sort of postprocessing after this run fn
-        if not synchronous:
-            os._exit(0)
-
         return submitted_run
 
 

wandb/sdk/launch/sweeps/__init__.py

@@ -3,9 +3,11 @@ from typing import Any, Callable, Dict
 
 log = logging.getLogger(__name__)
 
+SCHEDULER_URI = "placeholder-uri-scheduler"
+
 
 class SchedulerError(Exception):
-    """Raised when a known error occurs with wandb sweep scheduler"""
+    """Raised when a known error occurs with wandb sweep scheduler."""
 
     pass
 
@@ -22,7 +24,6 @@ _WANDB_SCHEDULERS: Dict[str, Callable] = {
 
 
 def load_scheduler(scheduler_name: str) -> Any:
-
     scheduler_name = scheduler_name.lower()
     if scheduler_name not in _WANDB_SCHEDULERS:
         raise SchedulerError(

wandb/sdk/launch/sweeps/scheduler.py

@@ -2,12 +2,14 @@
 import logging
 import os
 import threading
+import traceback
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import click
+import yaml
 
 import wandb
 import wandb.apis.public as public
@@ -16,21 +18,30 @@ from wandb.errors import CommError
 from wandb.sdk.launch.launch_add import launch_add
 from wandb.sdk.launch.sweeps import SchedulerError
 from wandb.sdk.lib.runid import generate_id
+from wandb.wandb_agent import Agent
 
-logger = logging.getLogger(__name__)
+_logger = logging.getLogger(__name__)
 LOG_PREFIX = f"{click.style('sched:', fg='cyan')} "
 
 
+@dataclass
+class _Worker:
+    agent_config: Dict[str, Any]
+    agent_id: str
+
+
 class SchedulerState(Enum):
     PENDING = 0
     STARTING = 1
     RUNNING = 2
-    COMPLETED = 3
-    FAILED = 4
-    STOPPED = 5
+    FLUSH_RUNS = 3
+    COMPLETED = 4
+    FAILED = 5
+    STOPPED = 6
+    CANCELLED = 7
 
 
-class SimpleRunState(Enum):
+class RunState(Enum):
     ALIVE = 0
     DEAD = 1
     UNKNOWN = 2
@@ -39,19 +50,16 @@ class SimpleRunState(Enum):
 @dataclass
 class SweepRun:
     id: str
-    state: SimpleRunState = SimpleRunState.ALIVE
+    state: RunState = RunState.ALIVE
     queued_run: Optional[public.QueuedRun] = None
     args: Optional[Dict[str, Any]] = None
     logs: Optional[List[str]] = None
-    program: Optional[str] = None
     # Threading can be used to run multiple workers in parallel
     worker_id: Optional[int] = None
 
 
 class Scheduler(ABC):
-    """The Scheduler is a controller/agent that will populate a Launch RunQueue
-    with jobs from a hyperparameter sweep.
-    """
+    """A controller/agent that populates a Launch RunQueue from a hyperparameter sweep."""
 
     def __init__(
         self,
@@ -73,18 +81,31 @@ class Scheduler(ABC):
         self._project = (
             project or os.environ.get("WANDB_PROJECT") or api.settings("project")
         )
+        self._sweep_id: str = sweep_id or "empty-sweep-id"
+        self._state: SchedulerState = SchedulerState.PENDING
+
         # Make sure the provided sweep_id corresponds to a valid sweep
         try:
-            self._api.sweep(sweep_id, "{}", entity=self._entity, project=self._project)
+            resp = self._api.sweep(
+                sweep_id, "{}", entity=self._entity, project=self._project
+            )
+            if resp.get("state") == SchedulerState.CANCELLED.name:
+                self._state = SchedulerState.CANCELLED
+            self._sweep_config = yaml.safe_load(resp["config"])
         except Exception as e:
             raise SchedulerError(f"{LOG_PREFIX}Exception when finding sweep: {e}")
-        self._sweep_id: str = sweep_id or "empty-sweep-id"
-        self._state: SchedulerState = SchedulerState.PENDING
+
         # Dictionary of the runs being managed by the scheduler
         self._runs: Dict[str, SweepRun] = {}
         # Threading lock to ensure thread-safe access to the runs dictionary
         self._threading_lock: threading.Lock = threading.Lock()
-        self._project_queue = project_queue or self._project
+        self._project_queue = project_queue
+        # Optionally run multiple workers in (pseudo-)parallel. Workers do not
+        # actually run training workloads, they simply send heartbeat messages
+        # (emulating a real agent) and add new runs to the launch queue. The
+        # launch agent is the one that actually runs the training workloads.
+        self._workers: Dict[int, _Worker] = {}
+
         # Scheduler may receive additional kwargs which will be piped into the launch command
         self._kwargs: Dict[str, Any] = kwargs
 
@@ -102,12 +123,12 @@ class Scheduler(ABC):
 
     @property
     def state(self) -> SchedulerState:
-        logger.debug(f"{LOG_PREFIX}Scheduler state is {self._state.name}")
+        _logger.debug(f"{LOG_PREFIX}Scheduler state is {self._state.name}")
         return self._state
 
     @state.setter
     def state(self, value: SchedulerState) -> None:
-        logger.debug(f"{LOG_PREFIX}Scheduler was {self.state.name} is {value.name}")
+        _logger.debug(f"{LOG_PREFIX}Scheduler was {self.state.name} is {value.name}")
         self._state = value
 
     def is_alive(self) -> bool:
@@ -115,17 +136,33 @@ class Scheduler(ABC):
             SchedulerState.COMPLETED,
             SchedulerState.FAILED,
             SchedulerState.STOPPED,
+            SchedulerState.CANCELLED,
         ]:
             return False
         return True
 
     def start(self) -> None:
+        """Start a scheduler, confirms prerequisites, begins execution loop."""
         wandb.termlog(f"{LOG_PREFIX}Scheduler starting.")
+        if not self.is_alive():
+            wandb.termerror(
+                f"{LOG_PREFIX}Sweep already {self.state.name.lower()}! Exiting..."
+            )
+            self.exit()
+            return
+
         self._state = SchedulerState.STARTING
+        if not self._try_load_executable():
+            wandb.termerror(
+                f"{LOG_PREFIX}No 'job' or 'image_uri' loaded from sweep config."
+            )
+            self.exit()
+            return
         self._start()
         self.run()
 
     def run(self) -> None:
+        """Main run function for all external schedulers."""
         wandb.termlog(f"{LOG_PREFIX}Scheduler Running.")
         self.state = SchedulerState.RUNNING
         try:
@@ -134,6 +171,11 @@ class Scheduler(ABC):
                     break
                 self._update_run_states()
                 self._run()
+                # if we hit the run_cap, now set to stopped after launching runs
+                if self.state == SchedulerState.FLUSH_RUNS:
+                    if len(self._runs.keys()) == 0:
+                        wandb.termlog(f"{LOG_PREFIX}Done polling on runs, exiting.")
+                        self.state = SchedulerState.STOPPED
         except KeyboardInterrupt:
             wandb.termlog(f"{LOG_PREFIX}Scheduler received KeyboardInterrupt. Exiting.")
             self.state = SchedulerState.STOPPED
@@ -157,6 +199,28 @@ class Scheduler(ABC):
             self.state = SchedulerState.FAILED
         self._stop_runs()
 
+    def _try_load_executable(self) -> bool:
+        """Check existance of valid executable for a run.
+
+        logs and returns False when job is unreachable
+        """
+        if self._kwargs.get("job"):
+            _public_api = public.Api()
+            try:
+                _job_artifact = _public_api.artifact(self._kwargs["job"], type="job")
+                wandb.termlog(
+                    f"{LOG_PREFIX}Successfully loaded job: {_job_artifact.name} in scheduler"
+                )
+            except Exception:
+                wandb.termerror(f"{LOG_PREFIX}{traceback.format_exc()}")
+                return False
+            return True
+        elif self._kwargs.get("image_uri"):
+            # TODO(gst): check docker existance? Use registry in launch config?
+            return True
+        else:
+            return False
+
     def _yield_runs(self) -> Iterator[Tuple[str, SweepRun]]:
         """Thread-safe way to iterate over the runs."""
         with self._threading_lock:
@@ -168,25 +232,38 @@ class Scheduler(ABC):
             self._stop_run(run_id)
 
     def _stop_run(self, run_id: str) -> None:
-        """Stops a run and removes it from the scheduler"""
+        """Stop a run and removes it from the scheduler."""
         if run_id in self._runs:
             run: SweepRun = self._runs[run_id]
-            run.state = SimpleRunState.DEAD
+            run.state = RunState.DEAD
             # TODO(hupo): Send command to backend to stop run
             wandb.termlog(f"{LOG_PREFIX} Stopped run {run_id}.")
 
     def _update_run_states(self) -> None:
+        """Iterate through runs.
+
+        Get state from backend and deletes runs if not in running state. Threadsafe.
+        """
         _runs_to_remove: List[str] = []
         for run_id, run in self._yield_runs():
             try:
                 _state = self._api.get_run_state(self._entity, self._project, run_id)
-                if _state is None or _state in [
-                    "crashed",
-                    "failed",
-                    "killed",
-                    "finished",
-                ]:
-                    run.state = SimpleRunState.DEAD
+                _rqi_state = run.queued_run.state if run.queued_run else None
+                if (
+                    not _state
+                    or _state
+                    in [
+                        "crashed",
+                        "failed",
+                        "killed",
+                        "finished",
+                    ]
+                    or _rqi_state == "failed"
+                ):
+                    _logger.debug(
+                        f"({run_id}) run-state:{_state}, rqi-state:{_rqi_state}"
+                    )
+                    run.state = RunState.DEAD
                     _runs_to_remove.append(run_id)
                 elif _state in [
                     "running",
@@ -194,12 +271,12 @@ class Scheduler(ABC):
                     "preempted",
                     "preempting",
                 ]:
-                    run.state = SimpleRunState.ALIVE
+                    run.state = RunState.ALIVE
             except CommError as e:
                 wandb.termlog(
                     f"{LOG_PREFIX}Issue when getting RunState for Run {run_id}: {e}"
                 )
-                run.state = SimpleRunState.UNKNOWN
+                run.state = RunState.UNKNOWN
                 continue
         # Remove any runs that are dead
         with self._threading_lock:
@@ -213,31 +290,47 @@ class Scheduler(ABC):
         entry_point: Optional[List[str]] = None,
         config: Optional[Dict[str, Any]] = None,
     ) -> "public.QueuedRun":
-        """Add a launch job to the Launch RunQueue."""
+        """Add a launch job to the Launch RunQueue.
+
+        run_id: supplied by gorilla from agentHeartbeat
+        entry_point: sweep entrypoint overrides image_uri/job entrypoint
+        config: launch config
+        """
+        # job and image first from CLI args, then from sweep config
+        _job = self._kwargs.get("job") or self._sweep_config.get("job")
+
+        _sweep_config_uri = self._sweep_config.get("image_uri")
+        _image_uri = self._kwargs.get("image_uri") or _sweep_config_uri
+        if _job is None and _image_uri is None:
+            raise SchedulerError(
+                f"{LOG_PREFIX}No 'job' nor 'image_uri' (run: {run_id})"
+            )
+        elif _job is not None and _image_uri is not None:
+            raise SchedulerError(f"{LOG_PREFIX}Sweep has both 'job' and 'image_uri'")
+
+        if self._sweep_config.get("command"):
+            entry_point = Agent._create_sweep_command(self._sweep_config["command"])
+            wandb.termwarn(
+                f"{LOG_PREFIX}Sweep command {entry_point} will override"
+                f' {"job" if _job else "image_uri"} entrypoint'
+            )
+
         run_id = run_id or generate_id()
-        # One of Job and URI is required
-        _job = self._kwargs.get("job", None)
-        _uri = self._kwargs.get("uri", None)
-        if _job is None and _uri is None:
-            # If no Job is specified, use a placeholder URI to prevent Launch failure
-            _uri = "placeholder-uri-queuedrun-from-scheduler"
-        # Queue is required
-        _queue = self._kwargs.get("queue", "default")
         queued_run = launch_add(
             run_id=run_id,
             entry_point=entry_point,
             config=config,
-            uri=_uri,
+            docker_image=_image_uri,  # TODO(gst): make agnostic (github? run uri?)
             job=_job,
             project=self._project,
             entity=self._entity,
-            queue_name=_queue,
+            queue_name=self._kwargs.get("queue"),
             project_queue=self._project_queue,
             resource=self._kwargs.get("resource", None),
             resource_args=self._kwargs.get("resource_args", None),
         )
         self._runs[run_id].queued_run = queued_run
         wandb.termlog(
-            f"{LOG_PREFIX}Added run to Launch RunQueue: {_queue} RunID:{run_id}."
+            f"{LOG_PREFIX}Added run to Launch queue: {self._kwargs.get('queue')} RunID:{run_id}."
         )
         return queued_run