Package a local directory code_source_path for AI Runtime tasks

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

@@ -0,0 +1 @@
+An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy` (honoring `.gitignore` and `sync.include`/`sync.exclude`, content-addressed so unchanged code skips re-upload), uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path. A `requirements.yaml` derived from the job's serverless `environments` is written next to the task's `command_path`. Complements the existing pre-built-tarball path by letting users point at a source directory directly.

acceptance/bin/list_code_snapshot.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""
+List the entries of the AI Runtime code snapshot tarball uploaded during deploy.
+
+Reads out.requests.txt, takes code_source_path from the jobs/create request,
+exports that workspace file via the CLI, and prints its sorted tar entries. Used
+to assert which local files the snapshot includes (gitignore / sync rules).
+"""
+
+import gzip
+import io
+import json
+import os
+import subprocess
+import sys
+import tarfile
+
+with open("out.requests.txt") as f:
+    text = f.read()
+
+# out.requests.txt is a stream of concatenated (pretty-printed) JSON objects.
+decoder = json.JSONDecoder()
+requests = []
+pos = 0
+while pos < len(text):
+    while pos < len(text) and text[pos].isspace():
+        pos += 1
+    if pos >= len(text):
+        break
+    obj, pos = decoder.raw_decode(text, pos)
+    requests.append(obj)
+
+code_source_path = None
+for req in requests:
+    body = req.get("body")
+    if isinstance(body, dict) and req.get("path", "").endswith("/jobs/create"):
+        code_source_path = body["tasks"][0]["ai_runtime_task"]["code_source_path"]
+
+if not code_source_path:
+    sys.exit("no jobs/create request with code_source_path in out.requests.txt")
+
+# code_source_path is an absolute workspace path (/Workspace/Users/.../files/...).
+remote = code_source_path
+cli = os.environ["CLI"]
+local = "code_snapshot.tar.gz"
+# MSYS_NO_PATHCONV stops Git Bash on Windows from rewriting the /Workspace path.
+env = {**os.environ, "MSYS_NO_PATHCONV": "1"}
+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)
+
+with tarfile.open(fileobj=io.BytesIO(data)) as tar:
+    for name in sorted(tar.getnames()):
+        print(name)

acceptance/bin/print_requests.py

@@ -226,9 +226,12 @@ def main():
 
     for req in filtered_requests:
         body = req.get("body")
-        if isinstance(body, dict):
-            for field in del_body_fields:
+        for field in del_body_fields:
+            # Drop the field from the parsed body, and (for non-JSON bodies) the
+            # top-level raw_body — used to strip binary/volatile upload payloads.
+            if isinstance(body, dict):
                 body.pop(field, None)
+            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/local_code_source/databricks.yml

@@ -0,0 +1,31 @@
+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:
+        - 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"
+            dependencies:
+              - torch>=2.0.0

acceptance/bundle/ai_runtime_task/local_code_source/output.txt

@@ -0,0 +1,105 @@
+
+=== deploy packages and uploads the local code source
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+=== the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded
+
+>>> list_code_snapshot.py
+src/.gitignore
+src/command.sh
+src/data/model.bin
+src/train.py
+
+=== code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written
+
+>>> print_requests.py --sort --del-body raw_body //.air_snapshots/ //jobs/create //files/src/requirements.yaml
+{
+  "method": "POST",
+  "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[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/src/requirements.yaml",
+  "q": {
+    "overwrite": "true"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml",
+  "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/src_[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"
+      }
+    ]
+  }
+}
+
+=== 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-body 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/src_[SNAPSHOT].tar.gz",
+  "q": {
+    "overwrite": "true"
+  }
+}

acceptance/bundle/ai_runtime_task/local_code_source/script

@@ -0,0 +1,21 @@
+# A local directory code_source_path is packaged into a content-addressed tarball
+# written into the bundle (.air_snapshots/), uploaded by normal bundle file sync,
+# and code_source_path/command_path rewritten; a requirements.yaml is written next
+# to command_path.
+
+title "deploy packages and uploads the local code source\n"
+trace $CLI bundle deploy
+
+title "the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded\n"
+trace list_code_snapshot.py
+
+title "code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written\n"
+# --del-body raw_body drops the binary tarball upload body (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-body raw_body '//.air_snapshots/' '//jobs/create' '//files/src/requirements.yaml'
+
+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-body raw_body '//.air_snapshots/'

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/test.toml

@@ -0,0 +1,12 @@
+RecordRequests = true
+
+Ignore = [
+    '.databricks',
+]
+
+# The archive is content-addressed: <dir>_<sha256[:16]>.tar.gz. The hash is stable
+# given the committed inputs, but collapse it to a token so the test does not pin a
+# specific digest.
+[[Repls]]
+Old = 'src_[0-9a-f]{16}\.tar\.gz'
+New = 'src_[SNAPSHOT].tar.gz'