File size: 8,927 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 |
"""Utilities for the agent."""
from typing import Any, Dict, Optional
import wandb
from wandb.apis.internal import Api
from wandb.docker import is_docker_installed
from wandb.sdk.launch.errors import LaunchError
from .builder.abstract import AbstractBuilder
from .environment.abstract import AbstractEnvironment
from .registry.abstract import AbstractRegistry
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 == "local":
from .environment.local_environment import LocalEnvironment
return LocalEnvironment.from_config(config)
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)
if env_type == "azure":
from .environment.azure_environment import AzureEnvironment
return AzureEnvironment.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.
wandb.termwarn(
"The `registry` block of the launch agent config is being deprecated. "
"Please specify an image repository URI under the `builder.destination` "
"key of your launch agent config. See "
"https://docs.wandb.ai/guides/launch/setup-agent-advanced#agent-configuration "
"for more information."
)
registry_type = config.get("type")
if registry_type is None or registry_type == "local":
from .registry.local_registry import LocalRegistry
return LocalRegistry() # This is the default, dummy registry.
if registry_type == "ecr":
from .registry.elastic_container_registry import ElasticContainerRegistry
return ElasticContainerRegistry.from_config(config)
if registry_type == "gcr":
from .registry.google_artifact_registry import GoogleArtifactRegistry
return GoogleArtifactRegistry.from_config(config)
if registry_type == "acr":
from .registry.azure_container_registry import AzureContainerRegistry
return AzureContainerRegistry.from_config(config)
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 default builder is returned.
The default builder will be a DockerBuilder if we find a working docker cli
on the system, otherwise it will be a NoOpBuilder.
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:
if is_docker_installed():
from .builder.docker_builder import DockerBuilder
return DockerBuilder.from_config(
{}, environment, registry
) # This is the default builder.
from .builder.noop import NoOpBuilder
return NoOpBuilder.from_config({}, environment, registry)
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":
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,
registry: AbstractRegistry,
) -> 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, registry)
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):
try:
environment = AwsEnvironment.from_default()
except LaunchError as e:
raise LaunchError(
"Could not create Sagemaker runner. "
"Environment must be an instance of AwsEnvironment."
) from e
from .runner.sagemaker_runner import SageMakerRunner
return SageMakerRunner(api, runner_config, environment, registry)
if runner_name in ["vertex", "gcp-vertex"]:
from .environment.gcp_environment import GcpEnvironment
if not isinstance(environment, GcpEnvironment):
try:
environment = GcpEnvironment.from_default()
except LaunchError as e:
raise LaunchError(
"Could not create Vertex runner. "
"Environment must be an instance of GcpEnvironment."
) from e
from .runner.vertex_runner import VertexRunner
return VertexRunner(api, runner_config, environment, registry)
if runner_name == "kubernetes":
from .runner.kubernetes_runner import KubernetesRunner
return KubernetesRunner(api, runner_config, environment, registry)
raise LaunchError(
f"Could not create runner from config. Invalid runner name: {runner_name}"
)
|