Package a local directory code_source_path for AI Runtime tasks

.nextchanges/bundles/ai-runtime-code-source-dir.md

@@ -0,0 +1 @@
+For jobs where `ai_runtime_task.code_source_path` is a relative path to a local directory, the directory is now packaged into a tarball (honoring `.gitignore` and `sync.include`/`sync.exclude`), uploaded during deployment, and `code_source_path` is rewritten to the uploaded workspace path.

acceptance/bin/list_code_snapshot.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""
+List the entries of each AI Runtime code snapshot tarball uploaded during deploy.
+
+Reads out.requests.txt, takes every ai_runtime_task.code_source_path from the
+jobs/create request, exports each workspace archive via the CLI, and prints its
+sorted tar entries (grouped per archive). Used to assert which local files each
+snapshot includes (gitignore / sync rules), across one or more tasks.
+"""
+
+import gzip
+import io
+import os
+import subprocess
+import sys
+import tarfile
+
+from print_requests import read_json_many
+
+
+def code_source_paths(requests):
+    """Every task's code_source_path from the jobs/create request(s)."""
+    result = []
+    for req in requests:
+        body = req.get("body")
+        if isinstance(body, dict) and req.get("path", "").endswith("/jobs/create"):
+            for task in body.get("tasks", []):
+                art = task.get("ai_runtime_task")
+                if art and art.get("code_source_path"):
+                    result.append(art["code_source_path"])
+    return result
+
+
+def print_entries(cli, env, remote):
+    local = "code_snapshot.tar.gz"
+    subprocess.run(
+        [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local],
+        check=True,
+        env=env,
+    )
+    with open(local, "rb") as f:
+        data = gzip.decompress(f.read())
+    os.remove(local)
+
+    # Print the archive's sync-relative name (hash tokenized by test.toml repls) so
+    # multi-archive output is legible.
+    print(f"# {remote.split('/files/', 1)[-1]}")
+    with tarfile.open(fileobj=io.BytesIO(data)) as tar:
+        for name in sorted(tar.getnames()):
+            print(name)
+
+
+def main():
+    with open("out.requests.txt") as f:
+        requests = read_json_many(f.read())
+
+    paths = code_source_paths(requests)
+    if not paths:
+        sys.exit("no jobs/create request with code_source_path in out.requests.txt")
+
+    cli = os.environ["CLI"]
+    # MSYS_NO_PATHCONV stops Git Bash on Windows from rewriting the /Workspace path.
+    env = {**os.environ, "MSYS_NO_PATHCONV": "1"}
+
+    # code_source_path is an absolute workspace path (/Workspace/Users/.../files/...).
+    # Sort so multi-task output is deterministic.
+    for remote in sorted(paths):
+        print_entries(cli, env, remote)
+
+
+if __name__ == "__main__":
+    main()

acceptance/bin/print_requests.py

@@ -201,10 +201,20 @@ def main():
         "body fields that diverge between deployment engines, e.g. identity fields the "
         "terraform provider serializes into the body but the direct engine sends as query params.",
     )
+    parser.add_argument(
+        "--del-field",
+        action="append",
+        default=[],
+        metavar="FIELDS",
+        help="Comma-separated top-level request fields to delete (repeatable). Unlike "
+        "--del-body, which edits the parsed JSON body, this drops a field of the request "
+        "record itself, e.g. raw_body for a binary upload payload.",
+    )
     parser.add_argument("--fname", default="out.requests.txt")
     args = parser.parse_args()
 
     del_body_fields = [field for group in args.del_body for field in group.split(",")]
+    del_fields = [field for group in args.del_field for field in group.split(",")]
 
     test_tmp_dir = os.environ.get("TEST_TMP_DIR")
     if test_tmp_dir:
@@ -229,6 +239,8 @@ def main():
         if isinstance(body, dict):
             for field in del_body_fields:
                 body.pop(field, None)
+        for field in del_fields:
+            req.pop(field, None)
     if args.verbose:
         print(
             f"Read {len(data)} chars, {len(requests)} requests, {len(filtered_requests)} after filtering",

acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml

@@ -0,0 +1,29 @@
+bundle:
+  name: ai-runtime-empty
+
+# An unrelated force-include elsewhere in the bundle. sync.include is added to the
+# file list regardless of the scoped walk, so it must not make an all-filtered code
+# directory look non-empty.
+sync:
+  include:
+    - assets/*.bin
+
+resources:
+  jobs:
+    train:
+      name: "[${bundle.target}] AI Runtime training"
+      tasks:
+        - task_key: train
+          environment_key: default
+          ai_runtime_task:
+            experiment: my-training
+            code_source_path: ./src
+            deployments:
+              - command_path: src/command.sh
+                compute:
+                  accelerator_type: GPU_8xH100
+                  accelerator_count: 8
+      environments:
+        - environment_key: default
+          spec:
+            environment_version: "5"

acceptance/bundle/ai_runtime_task/empty_code_source/output.txt

@@ -0,0 +1,4 @@
+
+>>> [CLI] bundle deploy
+Error: code_source_path "./src" has no files to package (all excluded by .gitignore or sync.exclude, or the directory is empty)
+

acceptance/bundle/ai_runtime_task/empty_code_source/script

@@ -0,0 +1,5 @@
+# When code_source_path resolves to a directory whose contents are all filtered out
+# (here: a src/.gitignore of "*"), there is nothing to package. Deploy must fail
+# with an actionable message rather than shipping an empty code archive.
+# A sync.include elsewhere in the bundle must not defeat that guard (see databricks.yml).
+musterr trace $CLI bundle deploy

acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore

@@ -0,0 +1 @@
+*

acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh

@@ -0,0 +1,2 @@
+cd $CODE_SOURCE_PATH
+python train.py

acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py

@@ -0,0 +1 @@
+print("x")

acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml

@@ -0,0 +1,43 @@
+bundle:
+  name: ai-runtime-test
+
+sync:
+  # *.log excluded; data/*.bin force-included despite .gitignore.
+  exclude:
+    - "**/*.log"
+  include:
+    - src/data/*.bin
+
+resources:
+  jobs:
+    train:
+      name: "[${bundle.target}] AI Runtime training"
+      tasks:
+        # Two AI Runtime tasks with distinct local code dirs: each is packaged into
+        # its own content-addressed tarball, both under the repo root's .air_snapshots.
+        - task_key: train
+          environment_key: default
+          ai_runtime_task:
+            experiment: my-training
+            code_source_path: ./src
+            deployments:
+              - command_path: src/command.sh
+                compute:
+                  accelerator_type: GPU_8xH100
+                  accelerator_count: 8
+        - task_key: train2
+          environment_key: default
+          ai_runtime_task:
+            experiment: my-training-2
+            code_source_path: ./src2
+            deployments:
+              - command_path: src2/command.sh
+                compute:
+                  accelerator_type: GPU_8xH100
+                  accelerator_count: 8
+      environments:
+        - environment_key: default
+          spec:
+            environment_version: "5"
+            dependencies:
+              - torch>=2.0.0

acceptance/bundle/ai_runtime_task/local_code_source/output.txt

@@ -0,0 +1,135 @@
+
+=== deploy packages and uploads the local code sources
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+=== each task's tarball holds only synced files (both under the repo root .air_snapshots)
+
+>>> list_code_snapshot.py
+# .air_snapshots/[SNAPSHOT].tar.gz
+src2/command.sh
+src2/train.py
+# .air_snapshots/[SNAPSHOT].tar.gz
+src/.gitignore
+src/command.sh
+src/data/model.bin
+src/train.py
+
+=== both code_source_paths point into the bundle .air_snapshots, command_paths rewritten, deps on environments spec
+
+>>> print_requests.py --sort --del-field raw_body //.air_snapshots/ //jobs/create
+{
+  "method": "POST",
+  "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz",
+  "q": {
+    "overwrite": "true"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz",
+  "q": {
+    "overwrite": "true"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.2/jobs/create",
+  "body": {
+    "deployment": {
+      "kind": "BUNDLE",
+      "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/state/metadata.json"
+    },
+    "edit_mode": "UI_LOCKED",
+    "environments": [
+      {
+        "environment_key": "default",
+        "spec": {
+          "dependencies": [
+            "torch>=2.0.0"
+          ],
+          "environment_version": "5"
+        }
+      }
+    ],
+    "format": "MULTI_TASK",
+    "max_concurrent_runs": 1,
+    "name": "[default] AI Runtime training",
+    "queue": {
+      "enabled": true
+    },
+    "tasks": [
+      {
+        "ai_runtime_task": {
+          "code_source_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz",
+          "deployments": [
+            {
+              "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh",
+              "compute": {
+                "accelerator_count": 8,
+                "accelerator_type": "GPU_8xH100"
+              }
+            }
+          ],
+          "experiment": "my-training"
+        },
+        "environment_key": "default",
+        "task_key": "train"
+      },
+      {
+        "ai_runtime_task": {
+          "code_source_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz",
+          "deployments": [
+            {
+              "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src2/command.sh",
+              "compute": {
+                "accelerator_count": 8,
+                "accelerator_type": "GPU_8xH100"
+              }
+            }
+          ],
+          "experiment": "my-training-2"
+        },
+        "environment_key": "default",
+        "task_key": "train2"
+      }
+    ]
+  }
+}
+
+=== re-planning unchanged code is a no-op (no changes)
+
+>>> [CLI] bundle plan
+Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged
+
+=== editing a file changes the snapshot hash (content-addressed name changes)
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> print_requests.py --sort --del-field raw_body //.air_snapshots/
+{
+  "method": "POST",
+  "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz",
+  "q": {
+    "overwrite": "true"
+  }
+}
+
+=== destroy removes the deployed bundle (including the synced snapshots)
+
+>>> [CLI] bundle destroy --auto-approve
+The following resources will be deleted:
+  delete resources.jobs.train
+
+All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default
+
+Deleting files...
+Destroy complete!

acceptance/bundle/ai_runtime_task/local_code_source/script

@@ -0,0 +1,29 @@
+# Two AI Runtime tasks each with a local directory code_source_path. Each is
+# packaged into a content-addressed tarball written into the bundle
+# (.air_snapshots/, at the repo root) and uploaded by normal bundle file sync;
+# code_source_path/command_path are rewritten. Pip deps ride on the job's
+# environments[].spec.dependencies (no requirements.yaml is synthesized).
+
+title "deploy packages and uploads the local code sources\n"
+trace $CLI bundle deploy
+
+title "each task's tarball holds only synced files (both under the repo root .air_snapshots)\n"
+trace list_code_snapshot.py
+
+title "both code_source_paths point into the bundle .air_snapshots, command_paths rewritten, deps on environments spec\n"
+# --del-field raw_body drops the binary tarball upload payload (kept readable). Filters
+# use a leading // so Git Bash on Windows does not path-convert them. --keep is not
+# passed, so print_requests.py consumes out.requests.txt.
+trace print_requests.py --sort --del-field raw_body '//.air_snapshots/' '//jobs/create'
+
+title "re-planning unchanged code is a no-op (no changes)\n"
+trace $CLI bundle plan
+
+title "editing a file changes the snapshot hash (content-addressed name changes)\n"
+update_file.py src/train.py 'print("training")' 'print("training v2")'
+trace $CLI bundle deploy
+trace print_requests.py --sort --del-field raw_body '//.air_snapshots/'
+
+title "destroy removes the deployed bundle (including the synced snapshots)\n"
+trace $CLI bundle destroy --auto-approve
+rm out.requests.txt

acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore

@@ -0,0 +1,2 @@
+ignored_by_git.txt
+data/

acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh

@@ -0,0 +1,2 @@
+cd $CODE_SOURCE_PATH
+python train.py

acceptance/bundle/ai_runtime_task/local_code_source/src/data/model.bin

@@ -0,0 +1 @@
+weights

acceptance/bundle/ai_runtime_task/local_code_source/src/data/notes.txt

@@ -0,0 +1 @@
+scratch

acceptance/bundle/ai_runtime_task/local_code_source/src/debug.log

@@ -0,0 +1 @@
+log

acceptance/bundle/ai_runtime_task/local_code_source/src/ignored_by_git.txt

@@ -0,0 +1 @@
+kept

acceptance/bundle/ai_runtime_task/local_code_source/src/train.py

@@ -0,0 +1 @@
+print("training")

acceptance/bundle/ai_runtime_task/local_code_source/src2/command.sh

@@ -0,0 +1,2 @@
+cd $CODE_SOURCE_PATH
+python train.py

acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py

@@ -0,0 +1 @@
+print("train2")