Make a single GCS object usable as Data

docs/guides/data.md

@@ -31,6 +31,33 @@ process_data(Data("./my_dataset/"))
 process_data(Data("gs://my-bucket/training-set/"))
 ```
 
+A `Data` object that names one file resolves to a file path, not to a
+directory. This applies to local files and to single GCS objects:
+
+```python
+@kinetic.run(accelerator="cpu")
+def load_weights(weights_path):
+  import h5py
+
+  with h5py.File(weights_path) as f:
+    return list(f.keys())
+
+
+# Local file
+load_weights(Data("./weights.h5"))
+
+# Single GCS object: no trailing slash
+load_weights(Data("gs://my-bucket/checkpoints/weights.h5"))
+```
+
+:::{important}
+For a GCS URI, the trailing slash tells Kinetic which of the two you
+mean. `Data("gs://my-bucket/dataset/")` is a directory. Without the
+slash, the same URI names one object. Kinetic shows a warning at submit
+time if a URI has no trailing slash and its last segment has no file
+extension.
+:::
+
 `Data` works as a function argument, as a value inside a list/dict, and as
 a value in the `volumes={...}` decorator argument:
 
@@ -121,8 +148,16 @@ def read_config(config_path):
 
 
 read_config(Data("./config.json", fuse=True))
+
+# A single GCS object works the same way
+read_config(Data("gs://my-bucket/configs/model.json", fuse=True))
 ```
 
+GCS FUSE can mount directories only. For a single object, Kinetic thus
+mounts the parent directory of that object. Your function receives the
+path of the object in that mount. The mount also shows the other objects
+in the directory, but Kinetic reads no data from them.
+
 You can mix FUSE-mounted and downloaded data in the same job:
 
 ```python
@@ -264,6 +299,25 @@ For single files, the blob is stored at `{hash}/{filename}`. For
 directories, the full tree is preserved under `{hash}/`. The returned
 GCS URI always points to the hash prefix directory, not individual files.
 
+A GCS-hosted `Data` object does not use this pipeline. `upload_data()`
+returns the URI that you gave. An `is_dir=False` ref thus has one of two
+forms, and `_download_data()` in `remote_runner.py` must accept both:
+
+| Source                         | Ref `uri`                          | The URI names |
+| ------------------------------ | ---------------------------------- | ------------- |
+| Uploaded local file            | `gs://bucket/ns/data-cache/{hash}` | a directory   |
+| `Data("gs://bucket/dir/f.h5")` | `gs://bucket/dir/f.h5`             | the object    |
+
+Only the second form names a blob. The download thus first tries to get
+that object. If the bucket has no such object, the download lists the
+URI as a prefix. Both branches put the file into the target directory
+with its own name. `resolve_data_refs()` then gives your function that
+file path.
+
+If an `is_dir=False` ref matches no object and no prefix, the runner
+raises `FileNotFoundError`. The message contains the URI. The ref does
+not resolve to an empty directory.
+
 ### FUSE mount implementation
 
 GCS FUSE can only mount directories, not individual files. The system
@@ -287,3 +341,20 @@ volume. The `only-dir` mount option scopes the mount to a specific GCS
 prefix. For single files (`is_dir=False`), the parent directory is
 mounted. The pod receives a `gke-gcsfuse/volumes: "true"` annotation to
 trigger the GCS FUSE sidecar injection.
+
+**File selection in the mount:** `_resolve_fuse_single_file()` in
+`remote_runner.py` changes the mounted directory into a file path. It
+reads the last segment of the ref URI. Then it searches the mount for
+that name. The two ref forms above use different branches, but both
+branches are exact:
+
+- A GCS-native ref names the object. The parent mount shows this object
+  and the other objects in the directory, and the search finds the
+  correct one.
+- An uploaded file has a ref that names its hash directory, and the
+  search thus finds nothing. That directory contains only one object,
+  and that object is the file.
+
+If the mount contains more than one entry and the named object is not
+there, the runner raises `FileNotFoundError`. The runner does not select
+a different object.

kinetic/backend/execution_test.py

@@ -534,6 +534,42 @@ def test_fuse_data_arg_creates_auto_mount(self, mock_upload):
     self.assertTrue(spec["is_dir"])
     self.assertTrue(spec["read_only"])
 
+  @mock.patch("kinetic.backend.execution.storage.upload_data")
+  def test_fuse_gcs_object_keeps_its_own_uri(self, mock_upload):
+    """A GCS-native object is already file-level; nothing is appended.
+
+    ``build_gcs_fuse_volumes`` mounts the object's parent, and the pod
+    picks the object back out of it by name — both need the URI to keep
+    naming the object.
+    """
+    mock_upload.side_effect = lambda bucket, data, project: data.path
+    tmp = _make_temp_path(self)
+
+    fuse_data = Data("gs://bucket/datasets/weights.h5", fuse=True)
+    ctx = self._make_ctx(args=(fuse_data,))
+    _prepare_artifacts(ctx, str(tmp))
+
+    spec = ctx.fuse_volume_specs[0]
+    self.assertEqual(spec["gcs_uri"], "gs://bucket/datasets/weights.h5")
+    self.assertFalse(spec["is_dir"])
+
+  @mock.patch("kinetic.backend.execution.storage.upload_data")
+  def test_fuse_uploaded_file_uri_gains_the_filename(self, mock_upload):
+    """An uploaded file's URI is the hash dir, so the name is appended."""
+    mock_upload.return_value = "gs://bucket/ns/data-cache/abc123"
+    tmp = _make_temp_path(self)
+    config = tmp / "config.json"
+    config.write_text("{}")
+
+    ctx = self._make_ctx(args=(Data(str(config), fuse=True),))
+    _prepare_artifacts(ctx, str(tmp))
+
+    spec = ctx.fuse_volume_specs[0]
+    self.assertEqual(
+      spec["gcs_uri"], "gs://bucket/ns/data-cache/abc123/config.json"
+    )
+    self.assertFalse(spec["is_dir"])
+
   @mock.patch("kinetic.backend.execution.storage.upload_data")
   def test_mixed_fuse_and_non_fuse_volumes(self, mock_upload):
     mock_upload.return_value = "gs://bucket/hash/"

kinetic/backend/k8s_utils_test.py

@@ -682,6 +682,21 @@ def test_gcs_native_single_file(self):
     mount_opts = vols[0]["csi"]["volumeAttributes"]["mountOptions"]
     self.assertIn("only-dir=datasets/configs", mount_opts)
 
+  def test_bucket_root_single_file_mounts_the_whole_bucket(self):
+    """An object at the root has no parent prefix to scope the mount to."""
+    specs = [
+      {
+        "gcs_uri": "gs://bucket/model.json",
+        "mount_path": "/tmp/fuse-data/0",
+        "is_dir": False,
+        "read_only": True,
+      }
+    ]
+    _, vols, _ = build_gcs_fuse_volumes(specs)
+    mount_opts = vols[0]["csi"]["volumeAttributes"]["mountOptions"]
+    self.assertNotIn("only-dir", mount_opts)
+    self.assertEqual(vols[0]["csi"]["volumeAttributes"]["bucketName"], "bucket")
+
   def test_annotations_set(self):
     specs = [
       {

kinetic/runner/remote_runner.py

@@ -12,6 +12,7 @@
 import json
 import os
 import pickle
+import posixpath
 import shutil
 import subprocess
 import sys
@@ -890,21 +891,53 @@ def resolve_volumes(
     _download_data(ref, mount_path, storage_client)
 
 
-def _resolve_fuse_single_file(mount_path: str) -> str | None:
-  """Find the single data file inside a FUSE mount directory.
+def _gcs_blob_path(uri: str) -> str:
+  """The object path of a `gs://bucket/path` URI, without the bucket."""
+  stripped = uri[len("gs://") :] if uri.startswith("gs://") else uri
+  _, _, path = stripped.partition("/")
+  return path.strip("/")
 
-  GCS FUSE mounts directories, not individual files.  For single-file
-  data refs the mount is scoped to the hash directory containing the
-  file, so a flat listing is sufficient.
 
-  Returns the file path, or `None` if no data file is found.
+def _resolve_fuse_single_file(mount_path: str, uri: str) -> str | None:
+  """Find the data file of a single-file ref inside its FUSE mount.
+
+  GCS FUSE mounts directories, not individual objects, so a single-file
+  ref is always mounted through a parent directory.  Which directory
+  that is depends on where the data came from:
+
+  - A GCS-native object (`gs://bucket/dir/file.h5`) is mounted through
+    its own parent, which usually holds unrelated sibling objects. The
+    object name taken from *uri* picks the right one.
+  - An uploaded local file is mounted through its content-hash
+    directory, which holds exactly that one file. Its *uri* names the
+    hash directory rather than the file, so no entry matches and the
+    lone entry is the file.
+
+  Returns the file path, or `None` if the mount holds no data file.
+
+  Raises:
+      FileNotFoundError: If the object named by *uri* is absent from a
+          mount holding several entries, so no entry can be picked
+          without guessing.
   """
+  name = posixpath.basename(_gcs_blob_path(uri))
+  if name:
+    candidate = os.path.join(mount_path, name)
+    if os.path.exists(candidate):
+      return candidate
   try:
     entries = os.listdir(mount_path)
   except OSError:
     return None
-  if entries:
+  if len(entries) == 1:
     return os.path.join(mount_path, entries[0])
+  if entries:
+    raise FileNotFoundError(
+      f"FUSE mount {mount_path} for {uri} holds {len(entries)} entries "
+      f"and none is named {name!r}: {sorted(entries)[:10]}. Kinetic "
+      f"mounts a single object through its parent directory and then "
+      f"picks the object out by name; check that {uri} exists."
+    )
   return None
 
 
@@ -1149,7 +1182,7 @@ def _resolve_ref(obj):
       # For FUSE-mounted single files, resolve to the actual file path
       # rather than returning the mount directory.
       if obj.get("fuse") and not obj.get("is_dir"):
-        resolved = _resolve_fuse_single_file(obj["mount_path"])
+        resolved = _resolve_fuse_single_file(obj["mount_path"], obj["uri"])
         if resolved:
           return resolved
       return obj["mount_path"]
@@ -1305,7 +1338,12 @@ def _download_hf_data(
 def _download_data(
   ref: dict, target_dir: str, storage_client: storage.Client
 ) -> None:
-  """Download data from a GCS URI (or HF URI) to a local directory."""
+  """Download data from a GCS URI (or HF URI) to a local directory.
+
+  *target_dir* is always a directory: a single-object ref lands in it
+  under the object's own name, which is what ``resolve_data_refs`` then
+  hands the user function as a file path.
+  """
   os.makedirs(target_dir, exist_ok=True)
   uri = ref["uri"]
 
@@ -1322,6 +1360,15 @@ def _download_data(
   prefix = parts[1].rstrip("/") if len(parts) > 1 else ""
   bucket = storage_client.bucket(bucket_name)
 
+  # A single-file ref addresses the object itself when the data was
+  # already on GCS (`gs://bucket/dir/file.h5`), and the content-hash
+  # directory holding it when the file was uploaded from the client.
+  # Only the first form names a blob, so try that before listing.
+  is_dir = ref.get("is_dir", True)
+  if not is_dir and prefix and _download_object(bucket, prefix, target_dir):
+    logging.info("Downloaded 1 file from %s to %s", uri, target_dir)
+    return
+
   blobs = bucket.list_blobs(prefix=prefix + "/")
   total_downloaded = 0
   batch = []
@@ -1356,6 +1403,31 @@ def _download_data(
     logging.info(
       "Downloaded %d files from %s to %s", total_downloaded, uri, target_dir
     )
+  elif not is_dir:
+    # Neither an object nor a prefix: resolving this ref would hand the
+    # user function an empty directory, so fail with the URI instead.
+    raise FileNotFoundError(
+      f"No data found at {uri}: bucket {bucket_name!r} holds neither an "
+      f"object of that name nor any object under it. Check the URI; a "
+      f"Data path that names a directory needs a trailing slash."
+    )
+
+
+def _download_object(
+  bucket: storage.Bucket, blob_name: str, target_dir: str
+) -> bool:
+  """Download the object at *blob_name* into *target_dir*, if it exists.
+
+  The object keeps its own name inside *target_dir*.  Returns whether
+  an object was found, so callers can fall back to a prefix listing.
+  """
+  blob = bucket.blob(blob_name)
+  if not blob.exists():
+    return False
+  blob.download_to_filename(
+    os.path.join(target_dir, posixpath.basename(blob_name))
+  )
+  return True
 
 
 def _download_from_gcs(client, gcs_path, local_path):