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

Sign up to get free protection for your applications and to get access to all the features.
Files changed (228) hide show
  1. wandb/__init__.py +2 -3
  2. wandb/apis/__init__.py +1 -3
  3. wandb/apis/importers/__init__.py +4 -0
  4. wandb/apis/importers/base.py +312 -0
  5. wandb/apis/importers/mlflow.py +113 -0
  6. wandb/apis/internal.py +29 -2
  7. wandb/apis/normalize.py +6 -5
  8. wandb/apis/public.py +163 -180
  9. wandb/apis/reports/_templates.py +6 -12
  10. wandb/apis/reports/report.py +1 -1
  11. wandb/apis/reports/runset.py +1 -3
  12. wandb/apis/reports/util.py +12 -10
  13. wandb/beta/workflows.py +57 -34
  14. wandb/catboost/__init__.py +1 -2
  15. wandb/cli/cli.py +215 -133
  16. wandb/data_types.py +63 -56
  17. wandb/docker/__init__.py +78 -16
  18. wandb/docker/auth.py +21 -22
  19. wandb/env.py +0 -1
  20. wandb/errors/__init__.py +8 -116
  21. wandb/errors/term.py +1 -1
  22. wandb/fastai/__init__.py +1 -2
  23. wandb/filesync/dir_watcher.py +8 -5
  24. wandb/filesync/step_prepare.py +76 -75
  25. wandb/filesync/step_upload.py +1 -2
  26. wandb/integration/catboost/__init__.py +1 -3
  27. wandb/integration/catboost/catboost.py +8 -14
  28. wandb/integration/fastai/__init__.py +7 -13
  29. wandb/integration/gym/__init__.py +35 -4
  30. wandb/integration/keras/__init__.py +3 -3
  31. wandb/integration/keras/callbacks/metrics_logger.py +9 -8
  32. wandb/integration/keras/callbacks/model_checkpoint.py +9 -9
  33. wandb/integration/keras/callbacks/tables_builder.py +31 -19
  34. wandb/integration/kfp/kfp_patch.py +20 -17
  35. wandb/integration/kfp/wandb_logging.py +1 -2
  36. wandb/integration/lightgbm/__init__.py +21 -19
  37. wandb/integration/prodigy/prodigy.py +6 -7
  38. wandb/integration/sacred/__init__.py +9 -12
  39. wandb/integration/sagemaker/__init__.py +1 -3
  40. wandb/integration/sagemaker/auth.py +0 -1
  41. wandb/integration/sagemaker/config.py +1 -1
  42. wandb/integration/sagemaker/resources.py +1 -1
  43. wandb/integration/sb3/sb3.py +8 -4
  44. wandb/integration/tensorboard/__init__.py +1 -3
  45. wandb/integration/tensorboard/log.py +8 -8
  46. wandb/integration/tensorboard/monkeypatch.py +11 -9
  47. wandb/integration/tensorflow/__init__.py +1 -3
  48. wandb/integration/xgboost/__init__.py +4 -6
  49. wandb/integration/yolov8/__init__.py +7 -0
  50. wandb/integration/yolov8/yolov8.py +250 -0
  51. wandb/jupyter.py +31 -35
  52. wandb/lightgbm/__init__.py +1 -2
  53. wandb/old/settings.py +2 -2
  54. wandb/plot/bar.py +1 -2
  55. wandb/plot/confusion_matrix.py +1 -3
  56. wandb/plot/histogram.py +1 -2
  57. wandb/plot/line.py +1 -2
  58. wandb/plot/line_series.py +4 -4
  59. wandb/plot/pr_curve.py +17 -20
  60. wandb/plot/roc_curve.py +1 -3
  61. wandb/plot/scatter.py +1 -2
  62. wandb/proto/v3/wandb_server_pb2.py +85 -39
  63. wandb/proto/v3/wandb_telemetry_pb2.py +10 -10
  64. wandb/proto/v4/wandb_server_pb2.py +51 -39
  65. wandb/proto/v4/wandb_telemetry_pb2.py +10 -10
  66. wandb/sdk/__init__.py +1 -3
  67. wandb/sdk/backend/backend.py +1 -1
  68. wandb/sdk/data_types/_dtypes.py +38 -30
  69. wandb/sdk/data_types/base_types/json_metadata.py +1 -3
  70. wandb/sdk/data_types/base_types/media.py +17 -17
  71. wandb/sdk/data_types/base_types/wb_value.py +33 -26
  72. wandb/sdk/data_types/helper_types/bounding_boxes_2d.py +91 -125
  73. wandb/sdk/data_types/helper_types/classes.py +1 -1
  74. wandb/sdk/data_types/helper_types/image_mask.py +12 -12
  75. wandb/sdk/data_types/histogram.py +5 -4
  76. wandb/sdk/data_types/html.py +1 -2
  77. wandb/sdk/data_types/image.py +11 -11
  78. wandb/sdk/data_types/molecule.py +3 -6
  79. wandb/sdk/data_types/object_3d.py +1 -2
  80. wandb/sdk/data_types/plotly.py +1 -2
  81. wandb/sdk/data_types/saved_model.py +10 -8
  82. wandb/sdk/data_types/video.py +1 -1
  83. wandb/sdk/integration_utils/data_logging.py +5 -5
  84. wandb/sdk/interface/artifacts.py +288 -266
  85. wandb/sdk/interface/interface.py +2 -3
  86. wandb/sdk/interface/interface_grpc.py +1 -1
  87. wandb/sdk/interface/interface_queue.py +1 -1
  88. wandb/sdk/interface/interface_relay.py +1 -1
  89. wandb/sdk/interface/interface_shared.py +1 -2
  90. wandb/sdk/interface/interface_sock.py +1 -1
  91. wandb/sdk/interface/message_future.py +1 -1
  92. wandb/sdk/interface/message_future_poll.py +1 -1
  93. wandb/sdk/interface/router.py +1 -1
  94. wandb/sdk/interface/router_queue.py +1 -1
  95. wandb/sdk/interface/router_relay.py +1 -1
  96. wandb/sdk/interface/router_sock.py +1 -1
  97. wandb/sdk/interface/summary_record.py +1 -1
  98. wandb/sdk/internal/artifacts.py +1 -1
  99. wandb/sdk/internal/datastore.py +2 -3
  100. wandb/sdk/internal/file_pusher.py +5 -3
  101. wandb/sdk/internal/file_stream.py +22 -19
  102. wandb/sdk/internal/handler.py +5 -4
  103. wandb/sdk/internal/internal.py +1 -1
  104. wandb/sdk/internal/internal_api.py +115 -55
  105. wandb/sdk/internal/job_builder.py +1 -3
  106. wandb/sdk/internal/profiler.py +1 -1
  107. wandb/sdk/internal/progress.py +4 -6
  108. wandb/sdk/internal/sample.py +1 -3
  109. wandb/sdk/internal/sender.py +28 -16
  110. wandb/sdk/internal/settings_static.py +5 -5
  111. wandb/sdk/internal/system/assets/__init__.py +1 -0
  112. wandb/sdk/internal/system/assets/cpu.py +3 -9
  113. wandb/sdk/internal/system/assets/disk.py +2 -4
  114. wandb/sdk/internal/system/assets/gpu.py +6 -18
  115. wandb/sdk/internal/system/assets/gpu_apple.py +2 -4
  116. wandb/sdk/internal/system/assets/interfaces.py +50 -22
  117. wandb/sdk/internal/system/assets/ipu.py +1 -3
  118. wandb/sdk/internal/system/assets/memory.py +7 -13
  119. wandb/sdk/internal/system/assets/network.py +4 -8
  120. wandb/sdk/internal/system/assets/open_metrics.py +283 -0
  121. wandb/sdk/internal/system/assets/tpu.py +1 -4
  122. wandb/sdk/internal/system/assets/trainium.py +26 -14
  123. wandb/sdk/internal/system/system_info.py +2 -3
  124. wandb/sdk/internal/system/system_monitor.py +52 -20
  125. wandb/sdk/internal/tb_watcher.py +12 -13
  126. wandb/sdk/launch/_project_spec.py +54 -65
  127. wandb/sdk/launch/agent/agent.py +374 -90
  128. wandb/sdk/launch/builder/abstract.py +61 -7
  129. wandb/sdk/launch/builder/build.py +81 -110
  130. wandb/sdk/launch/builder/docker_builder.py +181 -0
  131. wandb/sdk/launch/builder/kaniko_builder.py +419 -0
  132. wandb/sdk/launch/builder/noop.py +31 -12
  133. wandb/sdk/launch/builder/templates/_wandb_bootstrap.py +70 -20
  134. wandb/sdk/launch/environment/abstract.py +28 -0
  135. wandb/sdk/launch/environment/aws_environment.py +276 -0
  136. wandb/sdk/launch/environment/gcp_environment.py +271 -0
  137. wandb/sdk/launch/environment/local_environment.py +65 -0
  138. wandb/sdk/launch/github_reference.py +3 -8
  139. wandb/sdk/launch/launch.py +38 -29
  140. wandb/sdk/launch/launch_add.py +6 -8
  141. wandb/sdk/launch/loader.py +230 -0
  142. wandb/sdk/launch/registry/abstract.py +54 -0
  143. wandb/sdk/launch/registry/elastic_container_registry.py +163 -0
  144. wandb/sdk/launch/registry/google_artifact_registry.py +203 -0
  145. wandb/sdk/launch/registry/local_registry.py +62 -0
  146. wandb/sdk/launch/runner/abstract.py +1 -16
  147. wandb/sdk/launch/runner/{kubernetes.py → kubernetes_runner.py} +83 -95
  148. wandb/sdk/launch/runner/local_container.py +46 -22
  149. wandb/sdk/launch/runner/local_process.py +1 -4
  150. wandb/sdk/launch/runner/{aws.py → sagemaker_runner.py} +53 -212
  151. wandb/sdk/launch/runner/{gcp_vertex.py → vertex_runner.py} +38 -55
  152. wandb/sdk/launch/sweeps/__init__.py +3 -2
  153. wandb/sdk/launch/sweeps/scheduler.py +132 -39
  154. wandb/sdk/launch/sweeps/scheduler_sweep.py +80 -89
  155. wandb/sdk/launch/utils.py +101 -30
  156. wandb/sdk/launch/wandb_reference.py +2 -7
  157. wandb/sdk/lib/_settings_toposort_generate.py +166 -0
  158. wandb/sdk/lib/_settings_toposort_generated.py +201 -0
  159. wandb/sdk/lib/apikey.py +2 -4
  160. wandb/sdk/lib/config_util.py +4 -1
  161. wandb/sdk/lib/console.py +1 -3
  162. wandb/sdk/lib/deprecate.py +3 -3
  163. wandb/sdk/lib/file_stream_utils.py +7 -5
  164. wandb/sdk/lib/filenames.py +1 -1
  165. wandb/sdk/lib/filesystem.py +61 -5
  166. wandb/sdk/lib/git.py +1 -3
  167. wandb/sdk/lib/import_hooks.py +4 -7
  168. wandb/sdk/lib/ipython.py +8 -5
  169. wandb/sdk/lib/lazyloader.py +1 -3
  170. wandb/sdk/lib/mailbox.py +14 -4
  171. wandb/sdk/lib/proto_util.py +10 -5
  172. wandb/sdk/lib/redirect.py +15 -22
  173. wandb/sdk/lib/reporting.py +1 -3
  174. wandb/sdk/lib/retry.py +4 -5
  175. wandb/sdk/lib/runid.py +1 -3
  176. wandb/sdk/lib/server.py +15 -9
  177. wandb/sdk/lib/sock_client.py +1 -1
  178. wandb/sdk/lib/sparkline.py +1 -1
  179. wandb/sdk/lib/wburls.py +1 -1
  180. wandb/sdk/service/port_file.py +1 -2
  181. wandb/sdk/service/service.py +36 -13
  182. wandb/sdk/service/service_base.py +12 -1
  183. wandb/sdk/verify/verify.py +5 -7
  184. wandb/sdk/wandb_artifacts.py +142 -177
  185. wandb/sdk/wandb_config.py +5 -8
  186. wandb/sdk/wandb_helper.py +1 -1
  187. wandb/sdk/wandb_init.py +24 -13
  188. wandb/sdk/wandb_login.py +9 -9
  189. wandb/sdk/wandb_manager.py +39 -4
  190. wandb/sdk/wandb_metric.py +2 -6
  191. wandb/sdk/wandb_require.py +4 -15
  192. wandb/sdk/wandb_require_helpers.py +1 -9
  193. wandb/sdk/wandb_run.py +95 -141
  194. wandb/sdk/wandb_save.py +1 -3
  195. wandb/sdk/wandb_settings.py +149 -54
  196. wandb/sdk/wandb_setup.py +66 -46
  197. wandb/sdk/wandb_summary.py +13 -10
  198. wandb/sdk/wandb_sweep.py +6 -7
  199. wandb/sdk/wandb_watch.py +1 -1
  200. wandb/sklearn/calculate/confusion_matrix.py +1 -1
  201. wandb/sklearn/calculate/learning_curve.py +1 -1
  202. wandb/sklearn/calculate/summary_metrics.py +1 -3
  203. wandb/sklearn/plot/__init__.py +1 -1
  204. wandb/sklearn/plot/classifier.py +27 -18
  205. wandb/sklearn/plot/clusterer.py +4 -5
  206. wandb/sklearn/plot/regressor.py +4 -4
  207. wandb/sklearn/plot/shared.py +2 -2
  208. wandb/sync/__init__.py +1 -3
  209. wandb/sync/sync.py +4 -5
  210. wandb/testing/relay.py +11 -10
  211. wandb/trigger.py +1 -1
  212. wandb/util.py +106 -81
  213. wandb/viz.py +4 -4
  214. wandb/wandb_agent.py +50 -50
  215. wandb/wandb_controller.py +2 -3
  216. wandb/wandb_run.py +1 -2
  217. wandb/wandb_torch.py +1 -1
  218. wandb/xgboost/__init__.py +1 -2
  219. {wandb-0.13.10.dist-info → wandb-0.14.0.dist-info}/METADATA +6 -2
  220. {wandb-0.13.10.dist-info → wandb-0.14.0.dist-info}/RECORD +224 -209
  221. {wandb-0.13.10.dist-info → wandb-0.14.0.dist-info}/WHEEL +1 -1
  222. wandb/sdk/launch/builder/docker.py +0 -80
  223. wandb/sdk/launch/builder/kaniko.py +0 -393
  224. wandb/sdk/launch/builder/loader.py +0 -32
  225. wandb/sdk/launch/runner/loader.py +0 -50
  226. {wandb-0.13.10.dist-info → wandb-0.14.0.dist-info}/LICENSE +0 -0
  227. {wandb-0.13.10.dist-info → wandb-0.14.0.dist-info}/entry_points.txt +0 -0
  228. {wandb-0.13.10.dist-info → wandb-0.14.0.dist-info}/top_level.txt +0 -0
@@ -1,3 +1,4 @@
1
+ import hashlib
1
2
  import json
2
3
  import logging
3
4
  import os
@@ -15,9 +16,12 @@ from six.moves import shlex_quote
15
16
  import wandb
16
17
  import wandb.docker as docker
17
18
  from wandb.apis.internal import Api
18
- from wandb.errors import DockerError, ExecutionError, LaunchError
19
+ from wandb.sdk.launch.loader import (
20
+ builder_from_config,
21
+ environment_from_config,
22
+ registry_from_config,
23
+ )
19
24
 
20
- from ...lib.git import GitRepo
21
25
  from .._project_spec import (
22
26
  EntryPoint,
23
27
  EntrypointDefaults,
@@ -25,9 +29,14 @@ from .._project_spec import (
25
29
  compute_command_args,
26
30
  fetch_and_validate_project,
27
31
  )
28
- from ..utils import LAUNCH_CONFIG_FILE, LOG_PREFIX, resolve_build_and_registry_config
32
+ from ..utils import (
33
+ LAUNCH_CONFIG_FILE,
34
+ LOG_PREFIX,
35
+ ExecutionError,
36
+ LaunchError,
37
+ resolve_build_and_registry_config,
38
+ )
29
39
  from .abstract import AbstractBuilder
30
- from .loader import load_builder
31
40
 
32
41
  _logger = logging.getLogger(__name__)
33
42
 
@@ -123,6 +132,7 @@ PIP_TEMPLATE = """
123
132
  RUN python -m venv /env
124
133
  # make sure we install into the env
125
134
  ENV PATH="/env/bin:$PATH"
135
+
126
136
  COPY {requirements_files} ./
127
137
  {buildx_optional_prefix} {pip_install}
128
138
  """
@@ -178,12 +188,13 @@ def get_current_python_version() -> Tuple[str, str]:
178
188
  def get_base_setup(
179
189
  launch_project: LaunchProject, py_version: str, py_major: str
180
190
  ) -> str:
181
- """Fill in the Dockerfile templates for stage 2 of build. CPU version is built on python, GPU
182
- version is built on nvidia:cuda"""
191
+ """Fill in the Dockerfile templates for stage 2 of build.
183
192
 
193
+ CPU version is built on python, GPU version is built on nvidia:cuda.
194
+ """
184
195
  python_base_image = f"python:{py_version}-buster"
185
- if launch_project.cuda:
186
- cuda_version = launch_project.cuda_version or DEFAULT_CUDA_VERSION
196
+ if launch_project.cuda_base_image:
197
+ _logger.info(f"Using cuda base image: {launch_project.cuda_base_image}")
187
198
  # cuda image doesn't come with python tooling
188
199
  if py_major == "2":
189
200
  python_packages = [
@@ -200,7 +211,7 @@ def get_base_setup(
200
211
  "python3-setuptools",
201
212
  ]
202
213
  base_setup = CUDA_SETUP_TEMPLATE.format(
203
- cuda_base_image=f"nvidia/cuda:{cuda_version}-runtime",
214
+ cuda_base_image=launch_project.cuda_base_image,
204
215
  python_packages=" \\\n".join(python_packages),
205
216
  py_version=py_version,
206
217
  )
@@ -214,7 +225,7 @@ def get_base_setup(
214
225
 
215
226
 
216
227
  def get_env_vars_dict(launch_project: LaunchProject, api: Api) -> Dict[str, str]:
217
- """Generates environment variables for the project.
228
+ """Generate environment variables for the project.
218
229
 
219
230
  Arguments:
220
231
  launch_project: LaunchProject to generate environment variables for.
@@ -380,55 +391,6 @@ def generate_dockerfile(
380
391
  return dockerfile_contents
381
392
 
382
393
 
383
- _inspected_images = {}
384
-
385
-
386
- def docker_image_exists(docker_image: str, should_raise: bool = False) -> bool:
387
- """Checks if a specific image is already available,
388
- optionally raising an exception"""
389
- _logger.info("Checking if base image exists...")
390
- try:
391
- data = docker.run(["docker", "image", "inspect", docker_image])
392
- # always true, since return stderr defaults to false
393
- assert isinstance(data, str)
394
- parsed = json.loads(data)[0]
395
- _inspected_images[docker_image] = parsed
396
- return True
397
- except (DockerError, ValueError) as e:
398
- if should_raise:
399
- raise e
400
- _logger.info("Base image not found. Generating new base image")
401
- return False
402
-
403
-
404
- def docker_image_inspect(docker_image: str) -> Dict[str, Any]:
405
- """Get the parsed json result of docker inspect image_name"""
406
- if _inspected_images.get(docker_image) is None:
407
- docker_image_exists(docker_image, True)
408
- return _inspected_images.get(docker_image, {})
409
-
410
-
411
- def pull_docker_image(docker_image: str) -> None:
412
- """Pulls the requested docker image"""
413
- if docker_image_exists(docker_image):
414
- # don't pull images if they exist already, eg if they are local images
415
- return
416
- try:
417
- docker.run(["docker", "pull", docker_image])
418
- except DockerError as e:
419
- raise LaunchError(f"Docker server returned error: {e}")
420
-
421
-
422
- def construct_gcp_image_uri(
423
- launch_project: LaunchProject,
424
- gcp_repo: str,
425
- gcp_project: str,
426
- gcp_registry: str,
427
- ) -> str:
428
- base_uri = launch_project.image_uri
429
- return "/".join([gcp_registry, gcp_project, gcp_repo, base_uri])
430
-
431
-
432
394
  def construct_gcp_registry_uri(
433
395
  gcp_repo: str, gcp_project: str, gcp_registry: str
434
396
  ) -> str:
@@ -462,29 +424,11 @@ def _parse_existing_requirements(launch_project: LaunchProject) -> str:
462
424
  return requirements_line
463
425
 
464
426
 
465
- def _get_docker_image_uri(name: Optional[str], work_dir: str, image_id: str) -> str:
466
- """
467
- Returns an appropriate Docker image URI for a project based on the git hash of the specified
468
- working directory.
469
- :param name: The URI of the Docker repository with which to tag the image. The
470
- repository URI is used as the prefix of the image URI.
471
- :param work_dir: Path to the working directory in which to search for a git commit hash
472
- """
473
- name = name.replace(" ", "-") if name else "wandb-launch"
474
- # Optionally include first 7 digits of git SHA in tag name, if available.
475
-
476
- git_commit = GitRepo(work_dir).last_commit
477
- version_string = (
478
- ":" + str(git_commit[:7]) + image_id if git_commit else ":" + image_id
479
- )
480
- return name + version_string
481
-
482
-
483
427
  def _create_docker_build_ctx(
484
428
  launch_project: LaunchProject,
485
429
  dockerfile_contents: str,
486
430
  ) -> str:
487
- """Creates build context temp dir containing Dockerfile and project code, returning path to temp dir."""
431
+ """Create a build context temp dir for a Dockerfile and project code."""
488
432
  directory = tempfile.mkdtemp()
489
433
  dst_path = os.path.join(directory, "src")
490
434
  assert launch_project.project_dir is not None
@@ -511,8 +455,12 @@ def _create_docker_build_ctx(
511
455
 
512
456
 
513
457
  def join(split_command: List[str]) -> str:
514
- """Return a shell-escaped string from *split_command*."""
515
- return " ".join(shlex.quote(arg) for arg in split_command)
458
+ """Return a shell-escaped string from *split_command*.
459
+
460
+ Also remove quotes from double quoted strings. Ex:
461
+ "'local container queue'" --> "local container queue"
462
+ """
463
+ return " ".join(shlex.quote(arg.replace("'", "")) for arg in split_command)
516
464
 
517
465
 
518
466
  def construct_builder_args(
@@ -521,7 +469,7 @@ def construct_builder_args(
521
469
  ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
522
470
  registry_config = None
523
471
  if launch_config is not None:
524
- build_config = launch_config.get("build")
472
+ build_config = launch_config.get("builder")
525
473
  registry_config = launch_config.get("registry")
526
474
 
527
475
  default_launch_config = None
@@ -539,16 +487,12 @@ def construct_builder_args(
539
487
  def build_image_with_builder(
540
488
  builder: AbstractBuilder,
541
489
  launch_project: LaunchProject,
542
- repository: Optional[Any],
543
490
  entry_point: EntryPoint,
544
491
  ) -> Optional[str]:
545
- """
546
- Helper for testing and logging
547
- """
492
+ """Build image with testing and logging."""
548
493
  wandb.termlog(f"{LOG_PREFIX}Building docker image from uri source")
549
494
  image_uri: Optional[str] = builder.build_image(
550
495
  launch_project,
551
- repository,
552
496
  entry_point,
553
497
  )
554
498
  return image_uri
@@ -557,42 +501,69 @@ def build_image_with_builder(
557
501
  def build_image_from_project(
558
502
  launch_project: LaunchProject,
559
503
  api: Api,
560
- launch_config: Optional[Dict[str, Any]] = None,
561
- default_builder_type: Optional[str] = "docker",
504
+ launch_config: Dict[str, Any],
562
505
  ) -> str:
563
- """
564
- Accepts a reference to the Api class and a pre-computed launch_spec
565
- object, with an optional launch_config to set things like repository
566
- which is used in naming the output docker image, and build_type defaulting
567
- to docker (but could be used to build kube resource jobs w/ "kaniko")
506
+ """Construct a docker image from a project and returns the URI of the image.
568
507
 
569
- updates launch_project with the newly created docker image uri and
570
- returns the uri
508
+ Arguments:
509
+ launch_project: The project to build an image from.
510
+ api: The API object to use for fetching the project.
511
+ launch_config: The launch config to use for building the image.
512
+
513
+ Returns:
514
+ The URI of the built image.
571
515
  """
572
516
  assert launch_project.uri, "To build an image on queue a URI must be set."
517
+ launch_config = launch_config or {}
518
+ env_config = launch_config.get("environment", {})
519
+ if not isinstance(env_config, dict):
520
+ wrong_type = type(env_config).__name__
521
+ raise LaunchError(
522
+ f"Invalid environment config: {env_config} of type {wrong_type} "
523
+ "loaded from launch config. Expected dict."
524
+ )
525
+ environment = environment_from_config(env_config)
526
+
527
+ registry_config = launch_config.get("registry", {})
528
+ if not isinstance(registry_config, dict):
529
+ wrong_type = type(registry_config).__name__
530
+ raise LaunchError(
531
+ f"Invalid registry config: {registry_config} of type {wrong_type}"
532
+ " loaded from launch config. Expected dict."
533
+ )
534
+ registry = registry_from_config(registry_config, environment)
535
+
536
+ builder_config = launch_config.get("builder", {})
537
+ if not isinstance(builder_config, dict):
538
+ wrong_type = type(builder_config).__name__
539
+ raise LaunchError(
540
+ f"Invalid builder config: {builder_config} of type {wrong_type} "
541
+ "loaded from launch config. Expected dict."
542
+ )
543
+ builder = builder_from_config(builder_config, environment, registry)
573
544
 
574
- builder_config, registry_config = construct_builder_args(launch_config)
575
- launch_project = fetch_and_validate_project(launch_project, api)
576
- # Currently support either url or repository keywords in registry
577
- repository = registry_config.get("url") or registry_config.get("repository")
545
+ if not builder:
546
+ raise LaunchError("Unable to build image. No builder found.")
578
547
 
579
- if not builder_config.get("type"):
580
- wandb.termlog(f"{LOG_PREFIX}No builder found, defaulting to docker")
581
- builder_config["type"] = default_builder_type
548
+ launch_project = fetch_and_validate_project(launch_project, api)
582
549
 
583
- builder = load_builder(builder_config)
584
550
  entry_point: EntryPoint = launch_project.get_single_entry_point() or EntryPoint(
585
551
  name=EntrypointDefaults.PYTHON[-1],
586
552
  command=EntrypointDefaults.PYTHON,
587
553
  )
588
-
589
- image_uri = build_image_with_builder(
590
- builder,
591
- launch_project,
592
- repository,
593
- entry_point,
594
- )
554
+ wandb.termlog(f"{LOG_PREFIX}Building docker image from uri source")
555
+ image_uri = builder.build_image(launch_project, entry_point)
595
556
  if not image_uri:
596
557
  raise LaunchError("Error building image uri")
597
558
  else:
598
559
  return image_uri
560
+
561
+
562
+ def image_tag_from_dockerfile_and_source(
563
+ launch_project: LaunchProject, dockerfile_contents: str
564
+ ) -> str:
565
+ """Hashes the source and dockerfile contents into a unique tag."""
566
+ image_source_string = launch_project.get_image_source_string()
567
+ unique_id_string = image_source_string + dockerfile_contents
568
+ image_tag = hashlib.sha256(unique_id_string.encode("utf-8")).hexdigest()[:8]
569
+ return image_tag
@@ -0,0 +1,181 @@
1
+ """Implementation of the docker builder."""
2
+ import logging
3
+ import os
4
+ from typing import Any, Dict
5
+
6
+ import wandb
7
+ import wandb.docker as docker
8
+ from wandb.sdk.launch.builder.abstract import AbstractBuilder
9
+ from wandb.sdk.launch.environment.abstract import AbstractEnvironment
10
+ from wandb.sdk.launch.registry.abstract import AbstractRegistry
11
+
12
+ from .._project_spec import (
13
+ EntryPoint,
14
+ LaunchProject,
15
+ create_metadata_file,
16
+ get_entry_point_command,
17
+ )
18
+ from ..registry.local_registry import LocalRegistry
19
+ from ..utils import (
20
+ LOG_PREFIX,
21
+ LaunchDockerError,
22
+ LaunchError,
23
+ sanitize_wandb_api_key,
24
+ warn_failed_packages_from_build_logs,
25
+ )
26
+ from .build import (
27
+ _create_docker_build_ctx,
28
+ generate_dockerfile,
29
+ image_tag_from_dockerfile_and_source,
30
+ validate_docker_installation,
31
+ )
32
+
33
+ _GENERATED_DOCKERFILE_NAME = "Dockerfile.wandb-autogenerated"
34
+ _logger = logging.getLogger(__name__)
35
+
36
+
37
+ class DockerBuilder(AbstractBuilder):
38
+ """Builds a docker image for a project.
39
+
40
+ Attributes:
41
+ builder_config (Dict[str, Any]): The builder config.
42
+
43
+ """
44
+
45
+ builder_type = "docker"
46
+ base_image = "python:3.8"
47
+ target_platform = "linux/amd64"
48
+
49
+ def __init__(
50
+ self,
51
+ environment: AbstractEnvironment,
52
+ registry: AbstractRegistry,
53
+ verify: bool = True,
54
+ login: bool = True,
55
+ ):
56
+ """Initialize a DockerBuilder.
57
+
58
+ Arguments:
59
+ environment (AbstractEnvironment): The environment to use.
60
+ registry (AbstractRegistry): The registry to use.
61
+ verify (bool, optional): Whether to verify the functionality of the builder.
62
+ login (bool, optional): Whether to login to the registry.
63
+
64
+ Raises:
65
+ LaunchError: If docker is not installed
66
+ """
67
+ self.environment = environment # Docker builder doesn't actually use this.
68
+ self.registry = registry
69
+ if verify:
70
+ self.verify()
71
+ if login:
72
+ self.login()
73
+
74
+ @classmethod
75
+ def from_config(
76
+ cls,
77
+ config: Dict[str, Any],
78
+ environment: AbstractEnvironment,
79
+ registry: AbstractRegistry,
80
+ verify: bool = True,
81
+ ) -> "DockerBuilder":
82
+ """Create a DockerBuilder from a config.
83
+
84
+ Arguments:
85
+ config (Dict[str, Any]): The config.
86
+ registry (AbstractRegistry): The registry to use.
87
+ verify (bool, optional): Whether to verify the functionality of the builder.
88
+ login (bool, optional): Whether to login to the registry.
89
+
90
+ Returns:
91
+ DockerBuilder: The DockerBuilder.
92
+ """
93
+ # TODO the config for the docker builder as of yet is empty
94
+ # but ultimately we should add things like target platform, base image, etc.
95
+ return cls(environment, registry)
96
+
97
+ def verify(self) -> None:
98
+ """Verify the builder."""
99
+ validate_docker_installation()
100
+
101
+ def login(self) -> None:
102
+ """Login to the registry."""
103
+ if isinstance(self.registry, LocalRegistry):
104
+ _logger.info(f"{LOG_PREFIX}No registry configured, skipping login.")
105
+ else:
106
+ username, password = self.registry.get_username_password()
107
+ docker.login(username, password, self.registry.uri)
108
+
109
+ def build_image(
110
+ self,
111
+ launch_project: LaunchProject,
112
+ entrypoint: EntryPoint,
113
+ ) -> str:
114
+ """Build the image for the given project.
115
+
116
+ Arguments:
117
+ launch_project (LaunchProject): The project to build.
118
+ entrypoint (EntryPoint): The entrypoint to use.
119
+ """
120
+ dockerfile_str = generate_dockerfile(
121
+ launch_project, entrypoint, launch_project.resource, "docker"
122
+ )
123
+
124
+ image_tag = image_tag_from_dockerfile_and_source(launch_project, dockerfile_str)
125
+
126
+ repository = None if not self.registry else self.registry.get_repo_uri()
127
+ # if repo is set, use the repo name as the image name
128
+ if repository:
129
+ image_uri = f"{repository}:{image_tag}"
130
+ # otherwise, base the image name off of the source
131
+ # which the launch_project checks in image_name
132
+ else:
133
+ image_uri = f"{launch_project.image_name}:{image_tag}"
134
+
135
+ if not launch_project.build_required() and self.registry.check_image_exists(
136
+ image_uri
137
+ ):
138
+ return image_uri
139
+
140
+ _logger.info(
141
+ f"image {image_uri} does not already exist in repository, building."
142
+ )
143
+
144
+ entry_cmd = get_entry_point_command(entrypoint, launch_project.override_args)
145
+
146
+ create_metadata_file(
147
+ launch_project,
148
+ image_uri,
149
+ sanitize_wandb_api_key(" ".join(entry_cmd)),
150
+ dockerfile_str,
151
+ )
152
+ build_ctx_path = _create_docker_build_ctx(launch_project, dockerfile_str)
153
+ dockerfile = os.path.join(build_ctx_path, _GENERATED_DOCKERFILE_NAME)
154
+ try:
155
+ output = docker.build(
156
+ tags=[image_uri], file=dockerfile, context_path=build_ctx_path
157
+ )
158
+ warn_failed_packages_from_build_logs(output, image_uri)
159
+
160
+ except docker.DockerError as e:
161
+ raise LaunchDockerError(f"Error communicating with docker client: {e}")
162
+
163
+ try:
164
+ os.remove(build_ctx_path)
165
+ except Exception:
166
+ _msg = f"{LOG_PREFIX}Temporary docker context file {build_ctx_path} was not deleted."
167
+ _logger.info(_msg)
168
+
169
+ if repository:
170
+ reg, tag = image_uri.split(":")
171
+ wandb.termlog(f"{LOG_PREFIX}Pushing image {image_uri}")
172
+ push_resp = docker.push(reg, tag)
173
+ if push_resp is None:
174
+ raise LaunchError("Failed to push image to repository")
175
+ elif (
176
+ launch_project.resource == "sagemaker"
177
+ and f"The push refers to repository [{repository}]" not in push_resp
178
+ ):
179
+ raise LaunchError(f"Unable to push image to ECR, response: {push_resp}")
180
+
181
+ return image_uri