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

wandb/sdk/launch/launch.py

@@ -6,18 +6,18 @@ import yaml
 
 import wandb
 from wandb.apis.internal import Api
-from wandb.errors import ExecutionError, LaunchError
 
+from . import loader
 from ._project_spec import create_project_from_spec, fetch_and_validate_project
 from .agent import LaunchAgent
-from .builder import loader as builder_loader
 from .builder.build import construct_builder_args
-from .runner import loader
 from .runner.abstract import AbstractRun
 from .utils import (
     LAUNCH_CONFIG_FILE,
     LAUNCH_DEFAULT_PROJECT,
     PROJECT_SYNCHRONOUS,
+    ExecutionError,
+    LaunchError,
     construct_launch_spec,
     validate_launch_spec_source,
 )
@@ -25,37 +25,56 @@ from .utils import (
 _logger = logging.getLogger(__name__)
 
 
-def resolve_agent_config(
+def resolve_agent_config(  # noqa: C901
     api: Api,
     entity: Optional[str],
     project: Optional[str],
     max_jobs: Optional[int],
     queues: Optional[Tuple[str]],
+    config: Optional[str],
 ) -> Tuple[Dict[str, Any], Api]:
+    """Resolve the agent config.
+
+    Arguments:
+        api (Api): The api.
+        entity (str): The entity.
+        project (str): The project.
+        max_jobs (int): The max number of jobs.
+        queues (Tuple[str]): The queues.
+        config (str): The config.
+
+    Returns:
+        Tuple[Dict[str, Any], Api]: The resolved config and api.
+    """
     defaults = {
         "entity": api.default_entity,
         "project": LAUNCH_DEFAULT_PROJECT,
         "max_jobs": 1,
+        "max_schedulers": 1,
         "queues": [],
         "api_key": api.api_key,
         "base_url": api.settings("base_url"),
         "registry": {},
-        "build": {},
+        "builder": {},
         "runner": {},
     }
     user_set_project = False
     resolved_config: Dict[str, Any] = defaults
-    if os.path.exists(os.path.expanduser(LAUNCH_CONFIG_FILE)):
-        config = {}
-        with open(os.path.expanduser(LAUNCH_CONFIG_FILE)) as f:
+    config_path = config or os.path.expanduser(LAUNCH_CONFIG_FILE)
+    if os.path.isfile(config_path):
+        launch_config = {}
+        with open(config_path) as f:
             try:
-                config = yaml.safe_load(f)
-                print(config)
+                launch_config = yaml.safe_load(f)
             except yaml.YAMLError as e:
                 raise LaunchError(f"Invalid launch agent config: {e}")
-        if config.get("project") is not None:
+        if launch_config.get("project") is not None:
             user_set_project = True
-        resolved_config.update(dict(config))
+        resolved_config.update(launch_config.items())
+    elif config is not None:
+        raise LaunchError(
+            f"Could not find use specified launch config file: {config_path}"
+        )
     if os.environ.get("WANDB_PROJECT") is not None:
         resolved_config.update({"project": os.environ.get("WANDB_PROJECT")})
         user_set_project = True
@@ -132,7 +151,6 @@ def _run(
     resource_args: Optional[Dict[str, Any]],
     launch_config: Optional[Dict[str, Any]],
     synchronous: Optional[bool],
-    cuda: Optional[bool],
     api: Api,
     run_id: Optional[str],
     repository: Optional[str],
@@ -152,7 +170,6 @@ def _run(
         parameters,
         resource_args,
         launch_config,
-        cuda,
         run_id,
         repository,
     )
@@ -164,20 +181,15 @@ def _run(
     runner_config: Dict[str, Any] = {}
     runner_config[PROJECT_SYNCHRONOUS] = synchronous
 
-    if repository:  # override existing registry with CLI arg
-        launch_config = launch_config or {}
-        registry = launch_config.get("registry", {})
-        registry["url"] = repository
-        launch_config["registry"] = registry
-
-    build_config, registry_config = construct_builder_args(
-        launch_config,
-    )
+    config = launch_config or {}
+    build_config, registry_config = construct_builder_args(config)
 
-    builder = builder_loader.load_builder(build_config)
-    backend = loader.load_backend(resource, api, runner_config)
+    environment = loader.environment_from_config(config.get("environment", {}))
+    registry = loader.registry_from_config(registry_config, environment)
+    builder = loader.builder_from_config(build_config, environment, registry)
+    backend = loader.runner_from_config(resource, api, runner_config, environment)
     if backend:
-        submitted_run = backend.run(launch_project, builder, registry_config)
+        submitted_run = backend.run(launch_project, builder)
         # this check will always pass, run is only optional in the agent case where
         # a run queue id is present on the backend config
         assert submitted_run
@@ -203,7 +215,6 @@ def run(
     docker_image: Optional[str] = None,
     config: Optional[Dict[str, Any]] = None,
     synchronous: Optional[bool] = True,
-    cuda: Optional[bool] = None,
     run_id: Optional[str] = None,
     repository: Optional[str] = None,
 ) -> AbstractRun:
@@ -233,7 +244,6 @@ def run(
         asynchronous runs launched via this method will be terminated. If
         ``synchronous`` is True and the run fails, the current process will
         error out as well.
-    cuda: Whether to build a CUDA-enabled docker image or not
     run_id: ID for the run (To ultimately replace the :name: field)
     repository: string name of repository path for remote registry
 
@@ -276,7 +286,6 @@ def run(
         resource_args=resource_args,
         launch_config=config,
         synchronous=synchronous,
-        cuda=cuda,
         api=api,
         run_id=run_id,
         repository=repository,

wandb/sdk/launch/launch_add.py

@@ -4,7 +4,6 @@ from typing import Any, Dict, List, Optional
 import wandb
 import wandb.apis.public as public
 from wandb.apis.internal import Api
-from wandb.errors import LaunchError
 from wandb.sdk.launch._project_spec import (
     compute_command_args,
     create_project_from_spec,
@@ -13,6 +12,7 @@ from wandb.sdk.launch.builder.build import build_image_from_project
 from wandb.sdk.launch.utils import (
     LAUNCH_DEFAULT_PROJECT,
     LOG_PREFIX,
+    LaunchError,
     construct_launch_spec,
     validate_launch_spec_source,
 )
@@ -44,7 +44,6 @@ def launch_add(
     params: Optional[Dict[str, Any]] = None,
     project_queue: Optional[str] = None,
     resource_args: Optional[Dict[str, Any]] = None,
-    cuda: Optional[bool] = None,
     run_id: Optional[str] = None,
     build: Optional[bool] = False,
     repository: Optional[str] = None,
@@ -69,7 +68,6 @@ def launch_add(
         the parameters used to run the original run.
     resource_args: Resource related arguments for launching runs onto a remote backend.
         Will be stored on the constructed launch config under ``resource_args``.
-    cuda: Whether to build a CUDA-enabled docker image or not
     run_id: optional string indicating the id of the launched run
     build: optional flag defaulting to false, requires queue to be set
         if build, an image is created, creates a job artifact, pushes a reference
@@ -116,7 +114,6 @@ def launch_add(
         params,
         project_queue,
         resource_args,
-        cuda,
         run_id=run_id,
         build=build,
         repository=repository,
@@ -139,7 +136,6 @@ def _launch_add(
     params: Optional[Dict[str, Any]],
     project_queue: Optional[str],
     resource_args: Optional[Dict[str, Any]] = None,
-    cuda: Optional[bool] = None,
     run_id: Optional[str] = None,
     build: Optional[bool] = False,
     repository: Optional[str] = None,
@@ -158,7 +154,6 @@ def _launch_add(
         params,
         resource_args,
         config,
-        cuda,
         run_id,
         repository,
     )
@@ -174,7 +169,7 @@ def _launch_add(
             launch_spec["job"] = None
 
         launch_project = create_project_from_spec(launch_spec, api)
-        docker_image_uri = build_image_from_project(launch_project, api, config)
+        docker_image_uri = build_image_from_project(launch_project, api, config or {})
         run = wandb.run or wandb.init(
             project=launch_spec["project"],
             entity=launch_spec["entity"],
@@ -207,7 +202,10 @@ def _launch_add(
         if updated_spec.get("resource"):
             launch_spec["resource"] = updated_spec.get("resource")
 
-    wandb.termlog(f"{LOG_PREFIX}Added run to queue {project_queue}/{queue_name}.")
+    if project_queue == LAUNCH_DEFAULT_PROJECT:
+        wandb.termlog(f"{LOG_PREFIX}Added run to queue {queue_name}.")
+    else:
+        wandb.termlog(f"{LOG_PREFIX}Added run to queue {project_queue}/{queue_name}.")
     wandb.termlog(f"{LOG_PREFIX}Launch spec:\n{pprint.pformat(launch_spec)}\n")
     public_api = public.Api()
     container_job = False

wandb/sdk/launch/loader.py

@@ -0,0 +1,230 @@
+"""Utilities for the agent."""
+from typing import Any, Dict, Optional
+
+from wandb.apis.internal import Api
+from wandb.sdk.launch.utils import LaunchError
+
+from .builder.abstract import AbstractBuilder
+from .environment.abstract import AbstractEnvironment
+from .registry.abstract import AbstractRegistry
+from .registry.local_registry import LocalRegistry
+from .runner.abstract import AbstractRunner
+
+WANDB_RUNNERS = {
+    "local-container",
+    "local-process",
+    "kubernetes",
+    "vertex",
+    "sagemaker",
+}
+
+
+def environment_from_config(config: Optional[Dict[str, Any]]) -> AbstractEnvironment:
+    """Create an environment from a config.
+
+    This helper function is used to create an environment from a config. The
+    config should have a "type" key that specifies the type of environment to
+    create. The remaining keys are passed to the environment's from_config
+    method. If the config is None or empty, a LocalEnvironment is returned.
+
+    Arguments:
+        config (Dict[str, Any]): The config.
+
+    Returns:
+        Environment: The environment constructed.
+    """
+    if not config:
+        from .environment.local_environment import LocalEnvironment
+
+        return LocalEnvironment()  # This is the default, dummy environment.
+    env_type = config.get("type")
+    if not env_type:
+        raise LaunchError(
+            "Could not create environment from config. Environment type not specified!"
+        )
+    if env_type == "aws":
+        from .environment.aws_environment import AwsEnvironment
+
+        return AwsEnvironment.from_config(config)
+    if env_type == "gcp":
+        from .environment.gcp_environment import GcpEnvironment
+
+        return GcpEnvironment.from_config(config)
+    raise LaunchError(
+        f"Could not create environment from config. Invalid type: {env_type}"
+    )
+
+
+def registry_from_config(
+    config: Optional[Dict[str, Any]], environment: AbstractEnvironment
+) -> AbstractRegistry:
+    """Create a registry from a config.
+
+    This helper function is used to create a registry from a config. The
+    config should have a "type" key that specifies the type of registry to
+    create. The remaining keys are passed to the registry's from_config
+    method. If the config is None or empty, a LocalRegistry is returned.
+
+    Arguments:
+        config (Dict[str, Any]): The registry config.
+        environment (Environment): The environment of the registry.
+
+    Returns:
+        The registry if config is not None, otherwise None.
+
+    Raises:
+        LaunchError: If the registry is not configured correctly.
+    """
+    if not config:
+        from .registry.local_registry import LocalRegistry
+
+        return LocalRegistry()  # This is the default, dummy registry.
+    registry_type = config.get("type")
+    if registry_type is None:
+        from .registry.local_registry import LocalRegistry
+
+        return LocalRegistry()  # This is the default, dummy registry.
+    if registry_type == "ecr":
+        from .environment.aws_environment import AwsEnvironment
+
+        if not isinstance(environment, AwsEnvironment):
+            raise LaunchError(
+                "Could not create ECR registry. "
+                "Environment must be an instance of AWSEnvironment."
+            )
+        from .registry.elastic_container_registry import ElasticContainerRegistry
+
+        return ElasticContainerRegistry.from_config(config, environment)
+    if registry_type == "gcr":
+        from .environment.gcp_environment import GcpEnvironment
+
+        if not isinstance(environment, GcpEnvironment):
+            raise LaunchError(
+                "Could not create GCR registry. "
+                "Environment must be an instance of GCPEnvironment."
+            )
+        from .registry.google_artifact_registry import GoogleArtifactRegistry
+
+        return GoogleArtifactRegistry.from_config(config, environment)
+    raise LaunchError(
+        f"Could not create registry from config. Invalid registry type: {registry_type}"
+    )
+
+
+def builder_from_config(
+    config: Optional[Dict[str, Any]],
+    environment: AbstractEnvironment,
+    registry: AbstractRegistry,
+) -> AbstractBuilder:
+    """Create a builder from a config.
+
+    This helper function is used to create a builder from a config. The
+    config should have a "type" key that specifies the type of builder to import
+    and create. The remaining keys are passed to the builder's from_config
+    method. If the config is None or empty, a DockerBuilder is returned.
+
+    Arguments:
+        config (Dict[str, Any]): The builder config.
+        registry (Registry): The registry of the builder.
+
+    Returns:
+        The builder.
+
+    Raises:
+        LaunchError: If the builder is not configured correctly.
+    """
+    if not config:
+        from .builder.docker_builder import DockerBuilder
+
+        return DockerBuilder.from_config(
+            {}, environment, registry
+        )  # This is the default builder.
+
+    builder_type = config.get("type")
+    if builder_type is None:
+        raise LaunchError(
+            "Could not create builder from config. Builder type not specified"
+        )
+    if builder_type == "docker":
+        from .builder.docker_builder import DockerBuilder
+
+        return DockerBuilder.from_config(config, environment, registry)
+    if builder_type == "kaniko":
+        if isinstance(registry, LocalRegistry):
+            raise LaunchError(
+                "Could not create Kaniko builder. "
+                "Registry must be a remote registry."
+            )
+        from .builder.kaniko_builder import KanikoBuilder
+
+        return KanikoBuilder.from_config(config, environment, registry)
+    if builder_type == "noop":
+        from .builder.noop import NoOpBuilder
+
+        return NoOpBuilder.from_config(config, environment, registry)
+    raise LaunchError(
+        f"Could not create builder from config. Invalid builder type: {builder_type}"
+    )
+
+
+def runner_from_config(
+    runner_name: str,
+    api: Api,
+    runner_config: Dict[str, Any],
+    environment: AbstractEnvironment,
+) -> AbstractRunner:
+    """Create a runner from a config.
+
+    This helper function is used to create a runner from a config. The
+    config should have a "type" key that specifies the type of runner to import
+    and create. The remaining keys are passed to the runner's from_config
+    method. If the config is None or empty, a LocalContainerRunner is returned.
+
+    Arguments:
+        runner_name (str): The name of the backend.
+        api (Api): The API.
+        runner_config (Dict[str, Any]): The backend config.
+
+    Returns:
+        The runner.
+
+    Raises:
+        LaunchError: If the runner is not configured correctly.
+    """
+    if not runner_name or runner_name in ["local-container", "local"]:
+        from .runner.local_container import LocalContainerRunner
+
+        return LocalContainerRunner(api, runner_config, environment)
+    if runner_name == "local-process":
+        from .runner.local_process import LocalProcessRunner
+
+        return LocalProcessRunner(api, runner_config)
+    if runner_name == "sagemaker":
+        from .environment.aws_environment import AwsEnvironment
+
+        if not isinstance(environment, AwsEnvironment):
+            raise LaunchError(
+                "Could not create Sagemaker runner. "
+                "Environment must be an instance of AwsEnvironment."
+            )
+        from .runner.sagemaker_runner import SageMakerRunner
+
+        return SageMakerRunner(api, runner_config, environment)
+    if runner_name in ["vertex", "gcp-vertex"]:
+        from .environment.gcp_environment import GcpEnvironment
+
+        if not isinstance(environment, GcpEnvironment):
+            raise LaunchError(
+                "Could not create Vertex runner. "
+                "Environment must be an instance of GcpEnvironment."
+            )
+        from .runner.vertex_runner import VertexRunner
+
+        return VertexRunner(api, runner_config, environment)
+    if runner_name == "kubernetes":
+        from .runner.kubernetes_runner import KubernetesRunner
+
+        return KubernetesRunner(api, runner_config, environment)
+    raise LaunchError(
+        f"Could not create runner from config. Invalid runner name: {runner_name}"
+    )

wandb/sdk/launch/registry/abstract.py

@@ -0,0 +1,54 @@
+"""Abstract base class for registries."""
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+from ..environment.abstract import AbstractEnvironment
+
+
+class AbstractRegistry(ABC):
+    """Abstract base class for registries."""
+
+    uri: str
+
+    @abstractmethod
+    def verify(self) -> None:
+        """Verify that the registry is configured correctly."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry.
+
+        Returns:
+            (str, str): The username and password.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_repo_uri(self) -> str:
+        """Get the URI for a repository.
+
+        Returns:
+            str: The URI.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def check_image_exists(self, image_uri: str) -> bool:
+        """Check if an image exists in the registry.
+
+        Arguments:
+            image_uri (str): The URI of the image.
+
+        Returns:
+            bool: True if the image exists.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_config(
+        cls, config: dict, environment: "AbstractEnvironment", verify: bool = True
+    ) -> "AbstractRegistry":
+        """Create a registry from a config."""
+        raise NotImplementedError

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -0,0 +1,163 @@
+"""Implementation of Elastic Container Registry class for wandb launch."""
+import base64
+import logging
+from typing import Dict, Tuple
+
+from wandb.sdk.launch.environment.aws_environment import AwsEnvironment
+from wandb.sdk.launch.utils import LaunchError
+from wandb.util import get_module
+
+from .abstract import AbstractRegistry
+
+botocore = get_module(
+    "botocore",
+    required="AWS environment requires botocore to be installed. Please install "
+    "it with `pip install wandb[launch]`.",
+)
+
+_logger = logging.getLogger(__name__)
+
+
+class ElasticContainerRegistry(AbstractRegistry):
+    """Elastic Container Registry class.
+
+    Attributes:
+        repo_name (str): The name of the repository.
+        environment (AwsEnvironment): The AWS environment.
+        uri (str): The uri of the repository.
+    """
+
+    repo_name: str
+    environment: AwsEnvironment
+    uri: str
+
+    def __init__(self, repo_name: str, environment: AwsEnvironment) -> None:
+        """Initialize the Elastic Container Registry.
+
+        Arguments:
+            repo_name (str): The name of the repository.
+            environment (AwsEnvironment): The AWS environment.
+
+        Raises:
+            LaunchError: If there is an error verifying the registry.
+        """
+        super().__init__()
+        _logger.info(
+            f"Initializing Elastic Container Registry with repotisory {repo_name}."
+        )
+        self.repo_name = repo_name
+        self.environment = environment
+        self.verify()
+
+    @classmethod
+    def from_config(  # type: ignore[override]
+        cls,
+        config: Dict,
+        environment: AwsEnvironment,
+        verify: bool = True,
+    ) -> "ElasticContainerRegistry":
+        """Create an Elastic Container Registry from a config.
+
+        Arguments:
+            config (dict): The config.
+            environment (AwsEnvironment): The AWS environment.
+
+        Returns:
+            ElasticContainerRegistry: The Elastic Container Registry.
+        """
+        if config.get("type") != "ecr":
+            raise LaunchError(
+                f"Could not create ElasticContainerRegistry from config. Expected type 'ecr' "
+                f"but got '{config.get('type')}'."
+            )
+        repository = config.get("repository")
+        if not repository:
+            raise LaunchError(
+                "Could not create ElasticContainerRegistry from config. 'repository' is required."
+            )
+        return cls(repository, environment)
+
+    def verify(self) -> None:
+        """Verify that the registry is accessible and the configured repo exists.
+
+        Raises:
+            RegistryError: If there is an error verifying the registry.
+        """
+        _logger.debug("Verifying Elastic Container Registry.")
+        try:
+            session = self.environment.get_session()
+            client = session.client("ecr")
+            response = client.describe_repositories(repositoryNames=[self.repo_name])
+            self.uri = response["repositories"][0]["repositoryUri"].split("/")[0]
+
+        except botocore.exceptions.ClientError as e:
+            code = e.response["Error"]["Code"]
+            msg = e.response["Error"]["Message"]
+            # TODO: Log the code and the message here?
+            raise LaunchError(
+                f"Error verifying Elastic Container Registry: {code} {msg}"
+            )
+
+    def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry.
+
+        Returns:
+            (str, str): The username and password.
+
+        Raises:
+            RegistryError: If there is an error getting the username and password.
+        """
+        _logger.debug("Getting username and password for Elastic Container Registry.")
+        try:
+            session = self.environment.get_session()
+            client = session.client("ecr")
+            response = client.get_authorization_token()
+            username, password = base64.standard_b64decode(
+                response["authorizationData"][0]["authorizationToken"]
+            ).split(b":")
+            return username.decode("utf-8"), password.decode("utf-8")
+
+        except botocore.exceptions.ClientError as e:
+            code = e.response["Error"]["Code"]
+            msg = e.response["Error"]["Message"]
+            # TODO: Log the code and the message here?
+            raise LaunchError(f"Error getting username and password: {code} {msg}")
+
+    def get_repo_uri(self) -> str:
+        """Get the uri of the repository.
+
+        Returns:
+            str: The uri of the repository.
+        """
+        return self.uri + "/" + self.repo_name
+
+    def check_image_exists(self, image_uri: str) -> bool:
+        """Check if the image tag exists.
+
+        Arguments:
+            image_uri (str): The full image_uri.
+
+        Returns:
+            bool: True if the image tag exists.
+        """
+        uri, tag = image_uri.split(":")
+        if uri != self.get_repo_uri():
+            raise LaunchError(
+                f"Image uri {image_uri} does not match Elastic Container Registry uri {self.get_repo_uri()}."
+            )
+
+        _logger.debug("Checking if image tag exists.")
+        try:
+            session = self.environment.get_session()
+            client = session.client("ecr")
+            response = client.describe_images(
+                repositoryName=self.repo_name, imageIds=[{"imageTag": tag}]
+            )
+            return len(response["imageDetails"]) > 0
+
+        except botocore.exceptions.ClientError as e:
+            code = e.response["Error"]["Code"]
+            if code == "ImageNotFoundException":
+                return False
+            msg = e.response["Error"]["Message"]
+            raise LaunchError(f"Error checking if image tag exists: {code} {msg}")