wandb 0.17.4__py3-none-any.whl → 0.17.6__py3-none-any.whl

wandb/sdk/interface/interface_shared.py

@@ -137,6 +137,7 @@ class InterfaceShared(InterfaceBase):
         check_version: Optional[pb.CheckVersionRequest] = None,
         log_artifact: Optional[pb.LogArtifactRequest] = None,
         download_artifact: Optional[pb.DownloadArtifactRequest] = None,
+        link_artifact: Optional[pb.LinkArtifactRequest] = None,
         defer: Optional[pb.DeferRequest] = None,
         attach: Optional[pb.AttachRequest] = None,
         server_info: Optional[pb.ServerInfoRequest] = None,
@@ -184,6 +185,8 @@ class InterfaceShared(InterfaceBase):
             request.log_artifact.CopyFrom(log_artifact)
         elif download_artifact:
             request.download_artifact.CopyFrom(download_artifact)
+        elif link_artifact:
+            request.link_artifact.CopyFrom(link_artifact)
         elif defer:
             request.defer.CopyFrom(defer)
         elif attach:
@@ -242,7 +245,6 @@ class InterfaceShared(InterfaceBase):
         request: Optional[pb.Request] = None,
         telemetry: Optional[tpb.TelemetryRecord] = None,
         preempting: Optional[pb.RunPreemptingRecord] = None,
-        link_artifact: Optional[pb.LinkArtifactRecord] = None,
         use_artifact: Optional[pb.UseArtifactRecord] = None,
         output: Optional[pb.OutputRecord] = None,
         output_raw: Optional[pb.OutputRawRecord] = None,
@@ -282,8 +284,6 @@ class InterfaceShared(InterfaceBase):
             record.metric.CopyFrom(metric)
         elif preempting:
             record.preempting.CopyFrom(preempting)
-        elif link_artifact:
-            record.link_artifact.CopyFrom(link_artifact)
         elif use_artifact:
             record.use_artifact.CopyFrom(use_artifact)
         elif output:
@@ -393,10 +393,6 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(files=files)
         self._publish(rec)
 
-    def _publish_link_artifact(self, link_artifact: pb.LinkArtifactRecord) -> Any:
-        rec = self._make_record(link_artifact=link_artifact)
-        self._publish(rec)
-
     def _publish_use_artifact(self, use_artifact: pb.UseArtifactRecord) -> Any:
         rec = self._make_record(use_artifact=use_artifact)
         self._publish(rec)
@@ -411,6 +407,12 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_request(download_artifact=download_artifact)
         return self._deliver_record(rec)
 
+    def _deliver_link_artifact(
+        self, link_artifact: pb.LinkArtifactRequest
+    ) -> MailboxHandle:
+        rec = self._make_request(link_artifact=link_artifact)
+        return self._deliver_record(rec)
+
     def _publish_artifact(self, proto_artifact: pb.ArtifactRecord) -> None:
         rec = self._make_record(artifact=proto_artifact)
         self._publish(rec)

wandb/sdk/internal/handler.py

@@ -230,7 +230,7 @@ class HandleManager:
     def handle_files(self, record: Record) -> None:
         self._dispatch_record(record)
 
-    def handle_link_artifact(self, record: Record) -> None:
+    def handle_request_link_artifact(self, record: Record) -> None:
         self._dispatch_record(record)
 
     def handle_use_artifact(self, record: Record) -> None:

wandb/sdk/internal/internal_api.py

@@ -232,14 +232,14 @@ class Api:
 
         # todo: remove these hacky hacks after settings refactor is complete
         #  keeping this code here to limit scope and so that it is easy to remove later
-        extra_http_headers = self.settings("_extra_http_headers") or json.loads(
+        self._extra_http_headers = self.settings("_extra_http_headers") or json.loads(
             self._environ.get("WANDB__EXTRA_HTTP_HEADERS", "{}")
         )
-        extra_http_headers.update(_thread_local_api_settings.headers or {})
+        self._extra_http_headers.update(_thread_local_api_settings.headers or {})
 
         auth = None
         if self.access_token is not None:
-            extra_http_headers["Authorization"] = f"Bearer {self.access_token}"
+            self._extra_http_headers["Authorization"] = f"Bearer {self.access_token}"
         elif _thread_local_api_settings.cookies is None:
             auth = ("api", self.api_key or "")
 
@@ -253,7 +253,7 @@ class Api:
                     "User-Agent": self.user_agent,
                     "X-WANDB-USERNAME": env.get_username(env=self._environ),
                     "X-WANDB-USER-EMAIL": env.get_user_email(env=self._environ),
-                    **extra_http_headers,
+                    **self._extra_http_headers,
                 },
                 use_json=True,
                 # this timeout won't apply when the DNS lookup fails. in that case, it will be 60s

wandb/sdk/internal/sender.py

@@ -1,5 +1,7 @@
 """sender."""
 
+import contextlib
+import gzip
 import json
 import logging
 import os
@@ -66,6 +68,7 @@ else:
 if TYPE_CHECKING:
     from wandb.proto.wandb_internal_pb2 import (
         ArtifactManifest,
+        ArtifactManifestEntry,
         ArtifactRecord,
         HttpResponse,
         LocalInfo,
@@ -105,22 +108,18 @@ def _framework_priority() -> Generator[Tuple[str, str], None, None]:
 
 def _manifest_json_from_proto(manifest: "ArtifactManifest") -> Dict:
     if manifest.version == 1:
-        contents = {
-            content.path: {
-                "digest": content.digest,
-                "birthArtifactID": content.birth_artifact_id
-                if content.birth_artifact_id
-                else None,
-                "ref": content.ref if content.ref else None,
-                "size": content.size if content.size is not None else None,
-                "local_path": content.local_path if content.local_path else None,
-                "skip_cache": content.skip_cache,
-                "extra": {
-                    extra.key: json.loads(extra.value_json) for extra in content.extra
-                },
+        if manifest.manifest_file_path:
+            contents = {}
+            with gzip.open(manifest.manifest_file_path, "rt") as f:
+                for line in f:
+                    entry_json = json.loads(line)
+                    path = entry_json.pop("path")
+                    contents[path] = entry_json
+        else:
+            contents = {
+                content.path: _manifest_entry_from_proto(content)
+                for content in manifest.contents
             }
-            for content in manifest.contents
-        }
     else:
         raise ValueError(f"unknown artifact manifest version: {manifest.version}")
 
@@ -135,6 +134,19 @@ def _manifest_json_from_proto(manifest: "ArtifactManifest") -> Dict:
     }
 
 
+def _manifest_entry_from_proto(entry: "ArtifactManifestEntry") -> Dict:
+    birth_artifact_id = entry.birth_artifact_id if entry.birth_artifact_id else None
+    return {
+        "digest": entry.digest,
+        "birthArtifactID": birth_artifact_id,
+        "ref": entry.ref if entry.ref else None,
+        "size": entry.size if entry.size is not None else None,
+        "local_path": entry.local_path if entry.local_path else None,
+        "skip_cache": entry.skip_cache,
+        "extra": {extra.key: json.loads(extra.value_json) for extra in entry.extra},
+    }
+
+
 class ResumeState:
     resumed: bool
     step: int
@@ -1473,8 +1485,13 @@ class SendManager:
         # tbrecord watching threads are handled by handler.py
         pass
 
-    def send_link_artifact(self, record: "Record") -> None:
-        link = record.link_artifact
+    def send_request_link_artifact(self, record: "Record") -> None:
+        if not (record.control.req_resp or record.control.mailbox_slot):
+            raise ValueError(
+                f"Expected either `req_resp` or `mailbox_slot`, got: {record.control!r}"
+            )
+        result = proto_util._result_from_record(record)
+        link = record.request.link_artifact
         client_id = link.client_id
         server_id = link.server_id
         portfolio_name = link.portfolio_name
@@ -1490,7 +1507,9 @@ class SendManager:
                     client_id, server_id, portfolio_name, entity, project, aliases
                 )
             except Exception as e:
+                result.response.log_artifact_response.error_message = f'error linking artifact to "{entity}/{project}/{portfolio_name}"; error: {e}'
                 logger.warning("Failed to link artifact to portfolio: %s", e)
+        self._respond_result(result)
 
     def send_use_artifact(self, record: "Record") -> None:
         """Pretend to send a used artifact.
@@ -1579,6 +1598,10 @@ class SendManager:
         )
 
         self._job_builder._handle_server_artifact(res, artifact)
+
+        if artifact.manifest.manifest_file_path:
+            with contextlib.suppress(FileNotFoundError):
+                os.remove(artifact.manifest.manifest_file_path)
         return res
 
     def send_alert(self, record: "Record") -> None:

wandb/sdk/launch/_launch.py

@@ -211,7 +211,9 @@ async def _launch(
     launch_project = LaunchProject.from_spec(launch_spec, api)
     launch_project.fetch_and_validate_project()
     entrypoint = launch_project.get_job_entry_point()
-    image_uri = launch_project.docker_image  # Either set by user or None.
+    image_uri = (
+        launch_project.docker_image or launch_project.job_base_image
+    )  # Either set by user or None.
 
     # construct runner config.
     runner_config: Dict[str, Any] = {}
@@ -224,7 +226,7 @@ async def _launch(
         await environment.verify()
     registry = loader.registry_from_config(registry_config, environment)
     builder = loader.builder_from_config(build_config, environment, registry)
-    if not launch_project.docker_image:
+    if not (launch_project.docker_image or launch_project.job_base_image):
         assert entrypoint
         image_uri = await builder.build_image(launch_project, entrypoint, None)
     backend = loader.runner_from_config(

wandb/sdk/launch/_project_spec.py

@@ -7,6 +7,7 @@ import enum
 import json
 import logging
 import os
+import shutil
 import tempfile
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
@@ -112,6 +113,9 @@ class LaunchProject:
         self.sweep_id = sweep_id
         self.author = launch_spec.get("author")
         self.python_version: Optional[str] = launch_spec.get("python_version")
+        self._job_dockerfile: Optional[str] = None
+        self._job_build_context: Optional[str] = None
+        self._job_base_image: Optional[str] = None
         self.accelerator_base_image: Optional[str] = resource_args_build.get(
             "accelerator", {}
         ).get("base_image") or resource_args_build.get("cuda", {}).get("base_image")
@@ -131,8 +135,6 @@ class LaunchProject:
         self._queue_name: Optional[str] = None
         self._queue_entity: Optional[str] = None
         self._run_queue_item_id: Optional[str] = None
-        self._job_dockerfile: Optional[str] = None
-        self._job_build_context: Optional[str] = None
 
     def init_source(self) -> None:
         if self.docker_image is not None:
@@ -146,6 +148,21 @@ class LaunchProject:
             self.project_dir = os.getcwd()
             self._entry_point = self.override_entrypoint
 
+    def change_project_dir(self, new_dir: str) -> None:
+        """Change the project directory to a new directory."""
+        # Copy the contents of the old project dir to the new project dir.
+        old_dir = self.project_dir
+        if old_dir is not None:
+            shutil.copytree(
+                old_dir,
+                new_dir,
+                symlinks=True,
+                dirs_exist_ok=True,
+                ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc", ".git"),
+            )
+            shutil.rmtree(old_dir)
+        self.project_dir = new_dir
+
     def init_git(self, git_info: Dict[str, str]) -> None:
         self.git_version = git_info.get("version")
         self.git_repo = git_info.get("repo")
@@ -212,14 +229,23 @@ class LaunchProject:
     def job_build_context(self) -> Optional[str]:
         return self._job_build_context
 
+    @property
+    def job_base_image(self) -> Optional[str]:
+        return self._job_base_image
+
     def set_job_dockerfile(self, dockerfile: str) -> None:
         self._job_dockerfile = dockerfile
 
     def set_job_build_context(self, build_context: str) -> None:
         self._job_build_context = build_context
 
+    def set_job_base_image(self, base_image: str) -> None:
+        self._job_base_image = base_image
+
     @property
     def image_name(self) -> str:
+        if self.job_base_image is not None:
+            return self.job_base_image
         if self.docker_image is not None:
             return self.docker_image
         elif self.uri is not None:
@@ -299,10 +325,8 @@ class LaunchProject:
 
     def build_required(self) -> bool:
         """Checks the source to see if a build is required."""
-        # since the image tag for images built from jobs
-        # is based on the job version index, which is immutable
-        # we don't need to build the image for a job if that tag
-        # already exists
+        if self.job_base_image is not None:
+            return False
         if self.source != LaunchSource.JOB:
             return True
         return False
@@ -316,7 +340,9 @@ class LaunchProject:
         Returns:
             Optional[str]: The Docker image or None if not specified.
         """
-        return self._docker_image
+        if self._docker_image:
+            return self._docker_image
+        return None
 
     @docker_image.setter
     def docker_image(self, value: str) -> None:
@@ -336,7 +362,7 @@ class LaunchProject:
         # assuming project only has 1 entry point, pull that out
         # tmp fn until we figure out if we want to support multiple entry points or not
         if not self._entry_point:
-            if not self.docker_image:
+            if not self.docker_image and not self.job_base_image:
                 raise LaunchError(
                     "Project must have at least one entry point unless docker image is specified."
                 )

wandb/sdk/launch/agent/agent.py

@@ -717,7 +717,7 @@ class LaunchAgent:
         _, build_config, registry_config = construct_agent_configs(
             default_config, override_build_config
         )
-        image_uri = project.docker_image
+        image_uri = project.docker_image or project.job_base_image
         entrypoint = project.get_job_entry_point()
         environment = loader.environment_from_config(
             default_config.get("environment", {})
@@ -727,7 +727,11 @@ class LaunchAgent:
         backend = loader.runner_from_config(
             resource, api, backend_config, environment, registry
         )
-        if not (project.docker_image or isinstance(backend, LocalProcessRunner)):
+        if not (
+            project.docker_image
+            or project.job_base_image
+            or isinstance(backend, LocalProcessRunner)
+        ):
             assert entrypoint is not None
             image_uri = await builder.build_image(project, entrypoint, job_tracker)
 

wandb/sdk/launch/agent/run_queue_item_file_saver.py

@@ -2,7 +2,7 @@
 
 import os
 import sys
-from typing import List, Optional, Union
+from typing import List, Optional
 
 import wandb
 
@@ -17,9 +17,7 @@ FileSubtypes = Literal["warning", "error"]
 class RunQueueItemFileSaver:
     def __init__(
         self,
-        agent_run: Optional[
-            Union["wandb.sdk.wandb_run.Run", "wandb.sdk.lib.RunDisabled"]
-        ],
+        agent_run: Optional["wandb.sdk.wandb_run.Run"],
         run_queue_item_id: str,
     ):
         self.run_queue_item_id = run_queue_item_id

wandb/sdk/launch/builder/build.py

@@ -201,7 +201,7 @@ def get_requirements_section(
     # If there is a requirements.txt at root of build context, use that.
     if (base_path / "src" / "requirements.txt").exists():
         requirements_files += ["src/requirements.txt"]
-        deps_install_line = "pip install -r requirements.txt"
+        deps_install_line = "pip install uv && uv pip install -r requirements.txt"
         with open(base_path / "src" / "requirements.txt") as f:
             requirements = f.readlines()
         if not any(["wandb" in r for r in requirements]):
@@ -237,7 +237,9 @@ def get_requirements_section(
                 with open(base_path / "src" / "requirements.txt", "w") as f:
                     f.write("\n".join(project_deps))
                 requirements_files += ["src/requirements.txt"]
-                deps_install_line = "pip install -r requirements.txt"
+                deps_install_line = (
+                    "pip install uv && uv pip install -r requirements.txt"
+                )
                 return PIP_TEMPLATE.format(
                     buildx_optional_prefix=prefix,
                     requirements_files=" ".join(requirements_files),

wandb/sdk/launch/builder/kaniko_builder.py

@@ -263,11 +263,17 @@ class KanikoBuilder(AbstractBuilder):
         repo_uri = await self.registry.get_repo_uri()
         image_uri = repo_uri + ":" + image_tag
 
-        if (
-            not launch_project.build_required()
-            and await self.registry.check_image_exists(image_uri)
-        ):
-            return image_uri
+        # The DOCKER_CONFIG_SECRET option is mutually exclusive with the
+        # registry classes, so we must skip the check for image existence in
+        # that case.
+        if not launch_project.build_required():
+            if DOCKER_CONFIG_SECRET:
+                wandb.termlog(
+                    f"Skipping check for existing image {image_uri} due to custom dockerconfig."
+                )
+            else:
+                if await self.registry.check_image_exists(image_uri):
+                    return image_uri
 
         _logger.info(f"Building image {image_uri}...")
         _, api_client = await get_kube_context_and_api_client(
@@ -286,7 +292,12 @@ class KanikoBuilder(AbstractBuilder):
         wandb.termlog(f"{LOG_PREFIX}Created kaniko job {build_job_name}")
 
         try:
-            if isinstance(self.registry, AzureContainerRegistry):
+            # DOCKER_CONFIG_SECRET is a user provided dockerconfigjson. Skip our
+            # dockerconfig handling if it's set.
+            if (
+                isinstance(self.registry, AzureContainerRegistry)
+                and not DOCKER_CONFIG_SECRET
+            ):
                 dockerfile_config_map = client.V1ConfigMap(
                     metadata=client.V1ObjectMeta(
                         name=f"docker-config-{build_job_name}"
@@ -344,7 +355,10 @@ class KanikoBuilder(AbstractBuilder):
         finally:
             wandb.termlog(f"{LOG_PREFIX}Cleaning up resources")
             try:
-                if isinstance(self.registry, AzureContainerRegistry):
+                if (
+                    isinstance(self.registry, AzureContainerRegistry)
+                    and not DOCKER_CONFIG_SECRET
+                ):
                     await core_v1.delete_namespaced_config_map(
                         f"docker-config-{build_job_name}", "wandb"
                     )
@@ -498,7 +512,10 @@ class KanikoBuilder(AbstractBuilder):
                     "readOnly": True,
                 }
             )
-        if isinstance(self.registry, AzureContainerRegistry):
+        if (
+            isinstance(self.registry, AzureContainerRegistry)
+            and not DOCKER_CONFIG_SECRET
+        ):
             # Add the docker config map
             volumes.append(
                 {
@@ -533,7 +550,11 @@ class KanikoBuilder(AbstractBuilder):
         # Apply the rest of our defaults
         pod_labels["wandb"] = "launch"
         # This annotation is required to enable azure workload identity.
-        if isinstance(self.registry, AzureContainerRegistry):
+        # Don't add this label if using a docker config secret for auth.
+        if (
+            isinstance(self.registry, AzureContainerRegistry)
+            and not DOCKER_CONFIG_SECRET
+        ):
             pod_labels["azure.workload.identity/use"] = "true"
         pod_spec["restartPolicy"] = pod_spec.get("restartPolicy", "Never")
         pod_spec["activeDeadlineSeconds"] = pod_spec.get(

wandb/sdk/launch/builder/templates/_wandb_bootstrap.py

@@ -39,12 +39,13 @@ def install_deps(
         deps (str[], None): The dependencies that failed to install
     """
     try:
+        subprocess.check_output(["pip", "install", "uv"], stderr=subprocess.STDOUT)
         # Include only uri if @ is present
         clean_deps = [d.split("@")[-1].strip() if "@" in d else d for d in deps]
         index_args = ["--extra-index-url", extra_index] if extra_index else []
         print("installing {}...".format(", ".join(clean_deps)))
         opts = opts or []
-        args = ["pip", "install"] + opts + clean_deps + index_args
+        args = ["uv", "pip", "install"] + opts + clean_deps + index_args
         sys.stdout.flush()
         subprocess.check_output(args, stderr=subprocess.STDOUT)
         return failed

wandb/sdk/launch/inputs/internal.py

@@ -11,12 +11,14 @@ import os
 import pathlib
 import shutil
 import tempfile
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 import wandb
 import wandb.data_types
 from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.inputs.schema import META_SCHEMA
 from wandb.sdk.wandb_run import Run
+from wandb.util import get_module
 
 from .files import config_path_is_valid, override_file
 
@@ -62,11 +64,13 @@ class JobInputArguments:
         self,
         include: Optional[List[str]] = None,
         exclude: Optional[List[str]] = None,
+        schema: Optional[dict] = None,
         file_path: Optional[str] = None,
         run_config: Optional[bool] = None,
     ):
         self.include = include
         self.exclude = exclude
+        self.schema = schema
         self.file_path = file_path
         self.run_config = run_config
 
@@ -121,15 +125,70 @@ def _publish_job_input(
         exclude_paths=[_split_on_unesc_dot(path) for path in input.exclude]
         if input.exclude
         else [],
+        input_schema=input.schema,
         run_config=input.run_config,
         file_path=input.file_path or "",
     )
 
 
+def _replace_refs_and_allofs(schema: dict, defs: Optional[dict]) -> dict:
+    """Recursively fix JSON schemas with common issues.
+
+    1. Replaces any instances of $ref with their associated definition in defs
+    2. Removes any "allOf" lists that only have one item, "lifting" the item up
+    See test_internal.py for examples
+    """
+    ret: Dict[str, Any] = {}
+    if "$ref" in schema and defs:
+        # Reference found, replace it with its definition
+        def_key = schema["$ref"].split("#/$defs/")[1]
+        # Also run recursive replacement in case a ref contains more refs
+        return _replace_refs_and_allofs(defs.pop(def_key), defs)
+    for key, val in schema.items():
+        if isinstance(val, dict):
+            # Step into dicts recursively
+            new_val_dict = _replace_refs_and_allofs(val, defs)
+            ret[key] = new_val_dict
+        elif isinstance(val, list):
+            # Step into each item in the list
+            new_val_list = []
+            for item in val:
+                if isinstance(item, dict):
+                    new_val_list.append(_replace_refs_and_allofs(item, defs))
+                else:
+                    new_val_list.append(item)
+            # Lift up allOf blocks with only one item
+            if (
+                key == "allOf"
+                and len(new_val_list) == 1
+                and isinstance(new_val_list[0], dict)
+            ):
+                ret.update(new_val_list[0])
+            else:
+                ret[key] = new_val_list
+        else:
+            # For anything else (str, int, etc) keep it as-is
+            ret[key] = val
+    return ret
+
+
+def _validate_schema(schema: dict) -> None:
+    jsonschema = get_module(
+        "jsonschema",
+        required="Setting job schema requires the jsonschema package. Please install it with `pip install 'wandb[launch]'`.",
+        lazy=False,
+    )
+    validator = jsonschema.Draft202012Validator(META_SCHEMA)
+    errs = sorted(validator.iter_errors(schema), key=str)
+    if errs:
+        wandb.termwarn(f"Schema includes unhandled or invalid configurations:\n{errs}")
+
+
 def handle_config_file_input(
     path: str,
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     """Declare an overridable configuration file for a launch job.
 
@@ -151,9 +210,24 @@ def handle_config_file_input(
         path,
         dest,
     )
+    if schema:
+        # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+        # or the BaseModel class itself (e.g. schema=MySchema)
+        if hasattr(schema, "model_json_schema") and callable(
+            schema.model_json_schema  # type: ignore
+        ):
+            schema = schema.model_json_schema()
+        if not isinstance(schema, dict):
+            raise LaunchError(
+                "schema must be a dict, Pydantic model instance, or Pydantic model class."
+            )
+        defs = schema.pop("$defs", None)
+        schema = _replace_refs_and_allofs(schema, defs)
+        _validate_schema(schema)
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,
+        schema=schema,
         file_path=path,
         run_config=False,
     )
@@ -165,7 +239,9 @@ def handle_config_file_input(
 
 
 def handle_run_config_input(
-    include: Optional[List[str]] = None, exclude: Optional[List[str]] = None
+    include: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     """Declare wandb.config as an overridable configuration for a launch job.
 
@@ -175,9 +251,24 @@ def handle_run_config_input(
     If there is no active run, the include and exclude paths are staged and sent
     when a run is created.
     """
+    if schema:
+        # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+        # or the BaseModel class itself (e.g. schema=MySchema)
+        if hasattr(schema, "model_json_schema") and callable(
+            schema.model_json_schema  # type: ignore
+        ):
+            schema = schema.model_json_schema()
+        if not isinstance(schema, dict):
+            raise LaunchError(
+                "schema must be a dict, Pydantic model instance, or Pydantic model class."
+            )
+        defs = schema.pop("$defs", None)
+        schema = _replace_refs_and_allofs(schema, defs)
+        _validate_schema(schema)
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,
+        schema=schema,
         run_config=True,
         file_path=None,
     )