chore: tweak CI and fix file_manager error

Signed-off-by: mkornfield <mkornfield@nvidia.com>

Author: mckornfield

.github/actions/setup-docker-buildx/action.yaml

@@ -1,17 +1,19 @@
 name: Set up Docker Buildx
 description: >
-  Relocates Docker storage to /mnt for additional disk headroom on
-  GitHub-hosted runners, then creates a buildx docker-container builder.
+  Optionally relocates Docker storage to /mnt when that path is on a separate
+  larger partition, then creates a buildx docker-container builder.
 
 inputs:
   builder_name:
     description: Name of the buildx builder instance
     required: false
     default: nmp-builder
   relocate_docker_root:
-    description: Move Docker data-root to /mnt before creating the builder
+    description: >
+      Docker relocation strategy. auto relocates only when /mnt is on a
+      separate device from /; always forces relocation to /mnt; never skips.
     required: false
-    default: "true"
+    default: auto
 
 runs:
   using: composite
@@ -24,10 +26,29 @@ runs:
         builder_name="${{ inputs.builder_name }}"
         relocate="${{ inputs.relocate_docker_root }}"
 
-        if [ "$relocate" = "true" ] && [ -d /mnt ]; then
-          echo "=== Disk before Docker relocation ==="
-          df -h / /mnt || true
+        echo "=== Disk layout ==="
+        df -h / /mnt 2>/dev/null || df -h /
+
+        should_relocate=false
+        if [ "$relocate" = "never" ]; then
+          echo "Docker relocation disabled (relocate_docker_root=never)"
+        elif [ ! -d /mnt ]; then
+          echo "Skipping Docker relocation: /mnt does not exist"
+        elif [ "$relocate" = "always" ]; then
+          should_relocate=true
+          echo "Docker relocation forced (relocate_docker_root=always)"
+        else
+          root_dev="$(df --output=source / | tail -1 | tr -d ' ')"
+          mnt_dev="$(df --output=source /mnt | tail -1 | tr -d ' ')"
+          if [ "$root_dev" != "$mnt_dev" ]; then
+            should_relocate=true
+            echo "Relocating Docker storage: /mnt is on ${mnt_dev}, / is on ${root_dev}"
+          else
+            echo "Skipping Docker relocation: /mnt shares device ${root_dev} with /"
+          fi
+        fi
 
+        if [ "$should_relocate" = "true" ]; then
           sudo systemctl stop docker.socket docker.service 2>/dev/null || sudo service docker stop || true
           sudo mkdir -p /mnt/docker
           if [ -d /var/lib/docker ] && [ ! -L /var/lib/docker ]; then
@@ -47,5 +68,5 @@ runs:
         docker buildx inspect --bootstrap
 
         echo "=== Disk after Buildx bootstrap ==="
-        df -h / /mnt || true
+        df -h / /mnt 2>/dev/null || df -h /
         docker system df || true

docker/Dockerfile.safe-synthesizer-tasks

@@ -82,6 +82,7 @@ LABEL org.opencontainers.image.title="Safe Synthesizer Tasks" \
 
 # Install the Safe Synthesizer engine/CUDA runtime before copying Platform
 # sources, so plugin-only edits do not invalidate the large dependency layers.
+COPY plugins/nemo-safe-synthesizer/constraints.txt /tmp/safe-synthesizer-constraints.txt
 RUN uv venv /opt/venv
 RUN printf '%s\n' \
         flashinfer-cubin \
@@ -123,6 +124,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --index-strategy unsafe-best-match \
         --torch-backend cu129 \
         --overrides /tmp/safe-synthesizer-overrides.txt \
+        --constraints /tmp/safe-synthesizer-constraints.txt \
         --excludes /tmp/exclude-flashinfer-torch-vllm.txt \
         --requirements /tmp/safe-synthesizer-runtime.txt
 
@@ -134,6 +136,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --index-strategy unsafe-best-match \
         --torch-backend cu129 \
         --overrides /tmp/safe-synthesizer-overrides.txt \
+        --constraints /tmp/safe-synthesizer-constraints.txt \
         --excludes /tmp/exclude-torch-vllm.txt \
         --requirements /tmp/safe-synthesizer-runtime.txt
 
@@ -145,6 +148,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --index-strategy unsafe-best-match \
         --torch-backend cu129 \
         --overrides /tmp/safe-synthesizer-overrides.txt \
+        --constraints /tmp/safe-synthesizer-constraints.txt \
         --excludes /tmp/exclude-vllm.txt \
         --requirements /tmp/safe-synthesizer-runtime.txt
 
@@ -156,6 +160,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --index-strategy unsafe-best-match \
         --torch-backend cu129 \
         --overrides /tmp/safe-synthesizer-overrides.txt \
+        --constraints /tmp/safe-synthesizer-constraints.txt \
         --requirements /tmp/safe-synthesizer-runtime.txt \
         --requirements /tmp/vllm-cu129.txt
 

packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/file_manager.py

@@ -6,6 +6,7 @@
 import tempfile
 from abc import abstractmethod
 from dataclasses import dataclass, field
+from inspect import isawaitable
 from pathlib import Path
 from typing import Protocol
 
@@ -163,7 +164,9 @@ async def _validate_storage(self) -> None:
         except FileNotFoundError:
             if self.ensure_fileset_exists:
                 logger.info(f"Creating new fileset: [{self.fileset_name}] in workspace [{self.workspace}]")
-                await self._fs._sdk.files.filesets.create(name=self.fileset_name, workspace=self.workspace)
+                result = self.sdk.files.filesets.create(name=self.fileset_name, workspace=self.workspace)
+                if isawaitable(result):
+                    await result
             else:
                 raise FileStorageDoesNotExist(
                     f"Fileset [{self.fileset_name}] in workspace [{self.workspace}] does not exist."

packages/nmp_common/tests/jobs/test_file_manager.py

@@ -133,8 +133,7 @@ def test_fileset_validate_storage_creates(fileset_manager, mock_fileset_fs, mock
     """Test validate_storage creates fileset when missing."""
     mock_fileset_fs._info.side_effect = FileNotFoundError("not found")
     fileset_manager.validate_storage()
-    # The async version uses _fs._sdk.files.filesets.create
-    mock_fileset_fs._sdk.files.filesets.create.assert_called_once()
+    mock_sdk.files.filesets.create.assert_called_once_with(name="job-results-jobid-123", workspace="default")
 
 
 def test_fileset_upload_file(tmp_path, fileset_manager, mock_fileset_fs):
@@ -148,6 +147,31 @@ def test_fileset_upload_file(tmp_path, fileset_manager, mock_fileset_fs):
     assert result == "default/job-results-jobid-123#remote/test.txt"
 
 
+async def test_async_fileset_validate_storage_creates(mock_async_nmp_sdk, mock_fileset_fs):
+    """Test async validate_storage creates fileset when missing."""
+    from unittest import mock
+
+    from nmp.common.entities import DEFAULT_WORKSPACE
+    from nmp.common.jobs.file_manager import AsyncFilesetFileManager
+
+    mock_fileset_fs._info.side_effect = FileNotFoundError("not found")
+
+    with mock.patch("nemo_platform_plugin.jobs.file_manager.FilesetFileSystem") as mock_fs_class:
+        mock_fs_class.return_value = mock_fileset_fs
+        async_manager = AsyncFilesetFileManager(
+            workspace=DEFAULT_WORKSPACE,
+            fileset_name="job-results-jobid-123",
+            sdk=mock_async_nmp_sdk,
+        )
+
+    await async_manager.validate_storage()
+
+    mock_async_nmp_sdk.files.filesets.create.assert_awaited_once_with(
+        name="job-results-jobid-123",
+        workspace="default",
+    )
+
+
 def test_fileset_upload_directory(tmp_path, fileset_manager, mock_fileset_fs):
     """Test uploading a directory."""
     test_dir = tmp_path / "mydir"

plugins/nemo-safe-synthesizer/constraints.txt

@@ -30,3 +30,4 @@ tornado>=6.5.5
 # Runtime compatibility constraints.
 boto3>=1.40.46,<1.40.62
 botocore>=1.40.46,<1.40.62
+gliner<0.2.27