wandb 0.17.0rc2__py3-none-win_amd64.whl → 0.17.1__py3-none-win_amd64.whl

wandb/sdk/launch/agent/agent.py

@@ -9,7 +9,9 @@ import time
 import traceback
 from dataclasses import dataclass
 from multiprocessing import Event
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import yaml
 
 import wandb
 from wandb.apis.internal import Api
@@ -18,11 +20,11 @@ from wandb.sdk.launch._launch_add import launch_add
 from wandb.sdk.launch.runner.local_container import LocalSubmittedRun
 from wandb.sdk.launch.runner.local_process import LocalProcessRunner
 from wandb.sdk.launch.sweeps.scheduler import Scheduler
+from wandb.sdk.launch.utils import LAUNCH_CONFIG_FILE, resolve_build_and_registry_config
 from wandb.sdk.lib import runid
 
 from .. import loader
 from .._project_spec import LaunchProject
-from ..builder.build import construct_agent_configs
 from ..errors import LaunchDockerError, LaunchError
 from ..utils import (
     LAUNCH_DEFAULT_PROJECT,
@@ -134,6 +136,31 @@ class InternalAgentLogger:
         _logger.debug(f"{LOG_PREFIX}{message}")
 
 
+def construct_agent_configs(
+    launch_config: Optional[Dict] = None,
+    build_config: Optional[Dict] = None,
+) -> Tuple[Optional[Dict[str, Any]], Dict[str, Any], Dict[str, Any]]:
+    registry_config = None
+    environment_config = None
+    if launch_config is not None:
+        build_config = launch_config.get("builder")
+        registry_config = launch_config.get("registry")
+
+    default_launch_config = None
+    if os.path.exists(os.path.expanduser(LAUNCH_CONFIG_FILE)):
+        with open(os.path.expanduser(LAUNCH_CONFIG_FILE)) as f:
+            default_launch_config = (
+                yaml.safe_load(f) or {}
+            )  # In case the config is empty, we want it to be {} instead of None.
+        environment_config = default_launch_config.get("environment")
+
+    build_config, registry_config = resolve_build_and_registry_config(
+        default_launch_config, build_config, registry_config
+    )
+
+    return environment_config, build_config, registry_config
+
+
 class LaunchAgent:
     """Launch agent class which polls run given run queues and launches runs for wandb launch."""
 
@@ -173,7 +200,7 @@ class LaunchAgent:
             config: Config dictionary for the agent.
         """
         self._entity = config["entity"]
-        self._project = config.get("project", LAUNCH_DEFAULT_PROJECT)
+        self._project = LAUNCH_DEFAULT_PROJECT
         self._api = api
         self._base_url = self._api.settings().get("base_url")
         self._ticks = 0
@@ -194,6 +221,7 @@ class LaunchAgent:
         self._stopped_run_timeout = config.get(
             "stopped_run_timeout", DEFAULT_STOPPED_RUN_TIMEOUT
         )
+        self._known_warnings: List[str] = []
 
         # Get agent version from env var if present, otherwise wandb version
         self.version: str = "wandb@" + wandb.__version__
@@ -277,6 +305,8 @@ class LaunchAgent:
 
     def _init_agent_run(self) -> None:
         # TODO: has it been long enough that all backends support agents?
+        self._wandb_run = None
+
         if self.gorilla_supports_agents:
             settings = wandb.Settings(silent=True, disable_git=True)
             self._wandb_run = wandb.init(
@@ -286,8 +316,6 @@ class LaunchAgent:
                 id=self._name,
                 job_type=HIDDEN_AGENT_RUN_TYPE,
             )
-        else:
-            self._wandb_run = None
 
     @property
     def thread_ids(self) -> List[int]:
@@ -339,10 +367,7 @@ class LaunchAgent:
         if self._name:
             output_str += f"{self._name} "
         if self.num_running_jobs < self._max_jobs:
-            output_str += "polling on "
-            if self._project != LAUNCH_DEFAULT_PROJECT:
-                output_str += f"project {self._project}, "
-            output_str += f"queues {','.join(self._queues)}, "
+            output_str += f"polling on queues {','.join(self._queues)}, "
         output_str += (
             f"running {self.num_running_jobs} out of a maximum of {self._max_jobs} jobs"
         )
@@ -434,7 +459,6 @@ class LaunchAgent:
             # We retry for 60 seconds with an exponential backoff in case
             # upsert run is taking a while.
             logs = None
-            start_time = time.time()
             interval = 1
             while True:
                 called_init = self._check_run_exists_and_inited(
@@ -443,7 +467,7 @@ class LaunchAgent:
                     job_and_run_status.run_id,
                     job_and_run_status.run_queue_item_id,
                 )
-                if called_init or time.time() - start_time > RUN_INFO_GRACE_PERIOD:
+                if called_init or interval > RUN_INFO_GRACE_PERIOD:
                     break
                 if not called_init:
                     # Fetch the logs now if we don't get run info on the
@@ -692,7 +716,7 @@ class LaunchAgent:
             default_config, override_build_config
         )
         image_uri = project.docker_image
-        entrypoint = project.get_single_entry_point()
+        entrypoint = project.get_job_entry_point()
         environment = loader.environment_from_config(
             default_config.get("environment", {})
         )
@@ -790,14 +814,30 @@ class LaunchAgent:
         known_error = False
         try:
             run = job_tracker.run
-            status = (await run.get_status()).state
+            status = await run.get_status()
+            state = status.state
+
+            for warning in status.messages:
+                if warning not in self._known_warnings:
+                    self._known_warnings.append(warning)
+                    success = self._api.update_run_queue_item_warning(
+                        job_tracker.run_queue_item_id,
+                        warning,
+                        "Kubernetes",
+                        [],
+                    )
+                    if not success:
+                        _logger.warning(
+                            f"Error adding warning {warning} to run queue item {job_tracker.run_queue_item_id}"
+                        )
+                        self._known_warnings.remove(warning)
 
-            if status == "preempted" and job_tracker.entity == self._entity:
+            if state == "preempted" and job_tracker.entity == self._entity:
                 config = launch_spec.copy()
                 config["run_id"] = job_tracker.run_id
                 config["_resume_count"] = config.get("_resume_count", 0) + 1
                 with self._jobs_lock:
-                    job_tracker.completed_status = status
+                    job_tracker.completed_status = state
                 if config["_resume_count"] > MAX_RESUME_COUNT:
                     wandb.termlog(
                         f"{LOG_PREFIX}Run {job_tracker.run_id} has already resumed {MAX_RESUME_COUNT} times."
@@ -819,10 +859,10 @@ class LaunchAgent:
                 )
                 return True
             # TODO change these statuses to an enum
-            if status in ["stopped", "failed", "finished", "preempted"]:
+            if state in ["stopped", "failed", "finished", "preempted"]:
                 if job_tracker.is_scheduler:
                     wandb.termlog(f"{LOG_PREFIX}Scheduler finished with ID: {run.id}")
-                    if status == "failed":
+                    if state == "failed":
                         # on fail, update sweep state. scheduler run_id should == sweep_id
                         try:
                             self._api.set_sweep_state(
@@ -836,7 +876,7 @@ class LaunchAgent:
                 else:
                     wandb.termlog(f"{LOG_PREFIX}Job finished with ID: {run.id}")
                 with self._jobs_lock:
-                    job_tracker.completed_status = status
+                    job_tracker.completed_status = state
                 return True
 
             return False

wandb/sdk/launch/agent/config.py

@@ -192,9 +192,6 @@ class AgentConfig(BaseModel):
         default=[],
         description="The queues to use for this agent.",
     )
-    project: Optional[str] = Field(
-        description="The W&B project to use for this agent.",
-    )
     entity: Optional[str] = Field(
         description="The W&B entity to use for this agent.",
     )

wandb/sdk/launch/builder/abstract.py

@@ -7,6 +7,12 @@ from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
 
 from .._project_spec import EntryPoint, LaunchProject
+from ..registry.anon import AnonynmousRegistry
+from ..utils import (
+    AZURE_CONTAINER_REGISTRY_URI_REGEX,
+    ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
+    GCP_ARTIFACT_REGISTRY_URI_REGEX,
+)
 
 if TYPE_CHECKING:
     from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
@@ -87,3 +93,64 @@ class AbstractBuilder(ABC):
             LaunchError: If the builder cannot be used to build images.
         """
         raise NotImplementedError
+
+
+def registry_from_uri(uri: str) -> AbstractRegistry:
+    """Create a registry helper object from a uri.
+
+    This function parses the URI and determines which supported registry it
+    belongs to. It then creates a registry helper object for that registry.
+    The supported remote registry types are:
+    - Azure Container Registry
+    - Google Container Registry
+    - AWS Elastic Container Registry
+
+    The format of the URI is as follows:
+    - Azure Container Registry: <registry-name>.azurecr.io/<repo-name>/<image-name>
+    - Google Container Registry: <location>-docker.pkg.dev/<project-id>/<repo-name>/<image-name>
+    - AWS Elastic Container Registry: <account-id>.dkr.ecr.<region>.amazonaws.com/<repo-name>/<image-name>
+
+    Our classification of the registry is based on the domain name. For example,
+    if the uri contains `.azurecr.io`, we classify it as an Azure
+    Container Registry. If the uri contains `.dkr.ecr`, we classify
+    it as an AWS Elastic Container Registry. If the uri contains
+    `-docker.pkg.dev`, we classify it as a Google Artifact Registry.
+
+    This function will attempt to load the approriate cloud helpers for the
+
+    `https://` prefix is optional for all of the above.
+
+    Arguments:
+        uri: The uri to create a registry from.
+
+    Returns:
+        The registry.
+
+    Raises:
+        LaunchError: If the registry helper cannot be loaded for the given URI.
+    """
+    if uri.startswith("https://"):
+        uri = uri[len("https://") :]
+
+    if AZURE_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.azure_container_registry import (
+            AzureContainerRegistry,
+        )
+
+        return AzureContainerRegistry(uri=uri)
+
+    elif GCP_ARTIFACT_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.google_artifact_registry import (
+            GoogleArtifactRegistry,
+        )
+
+        return GoogleArtifactRegistry(uri=uri)
+
+    elif ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(uri) is not None:
+        from wandb.sdk.launch.registry.elastic_container_registry import (
+            ElasticContainerRegistry,
+        )
+
+        return ElasticContainerRegistry(uri=uri)
+
+    return AnonynmousRegistry(uri=uri)