dh-cli 0.12.0__tar.gz → 0.12.2__tar.gz

{dh_cli-0.12.0 → dh_cli-0.12.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dh-cli
-Version: 0.12.0
+Version: 0.12.2
 Summary: Dayhoff Labs developer CLI
 Author-email: Dayhoff Labs <dev@dayhofflabs.com>
 License: # PolyForm Noncommercial License 1.0.0

{dh_cli-0.12.0 → dh_cli-0.12.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dh-cli"
-version = "0.12.0"
+version = "0.12.2"
 description = "Dayhoff Labs developer CLI"
 requires-python = ">=3.11"
 readme = "README.md"

{dh_cli-0.12.0 → dh_cli-0.12.2}/src/dh_cli/batch/commands/protmpnn.py

@@ -71,6 +71,11 @@ DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:protmp
     is_flag=True,
     help="Save per-residue log probabilities to scores/ directory",
 )
+@click.option(
+    "--image",
+    default=None,
+    help="Override the container image URI or digest (e.g. repo@sha256:...)",
+)
 def protmpnn(
     input_dir,
     workers,
@@ -83,6 +88,7 @@ def protmpnn(
     after,
     auto_validate_top,
     save_score,
+    image,
 ):
     """Design protein sequences with ProtMPNN/LigandMPNN.
 
@@ -136,7 +142,7 @@ def protmpnn(
             _run_local_mode(input_path, auto_validate_top, base_path, save_score)
             return
 
-    _submit_batch_job(input_path, workers, queue, dry_run, base_path, after, auto_validate_top, save_score)
+    _submit_batch_job(input_path, workers, queue, dry_run, base_path, after, auto_validate_top, save_score, image)
 
 
 def _has_local_gpu() -> bool:
@@ -184,6 +190,7 @@ def _submit_batch_job(
     after: tuple[str, ...] = (),
     auto_validate_top: int | None = None,
     save_score: bool = False,
+    image: str | None = None,
 ):
     """Submit ProtMPNN job to AWS Batch."""
     click.echo(f"Scanning {input_path} for YAML files...")
@@ -236,7 +243,7 @@ def _submit_batch_job(
         user=job_id.split("-")[0],
         pipeline="protmpnn",
         status=JobStatus.PENDING,
-        image_uri=DEFAULT_IMAGE_URI,
+        image_uri=image or DEFAULT_IMAGE_URI,
         input=InputConfig(
             source=str(input_path),
             num_sequences=num_files,
@@ -282,6 +289,7 @@ def _submit_batch_job(
             retry_attempts=5,
             depends_on=depends_on,
             share_identifier=get_aws_username(),
+            image_override=image,
         )
 
         manifest.status = JobStatus.SUBMITTED

{dh_cli-0.12.0 → dh_cli-0.12.2}/src/dh_cli/batch/commands/rfd3.py

@@ -95,13 +95,13 @@ def rfd3(
     input_path = Path(input_dir).resolve()
 
     if run_local:
-        _run_local_mode(input_path, base_path)
+        _run_local_mode(input_path, base_path, num_designs)
         return
 
     if not run_remote and not dry_run:
         if _has_local_gpu():
             click.echo("GPU detected — running locally (use --remote to override)")
-            _run_local_mode(input_path, base_path)
+            _run_local_mode(input_path, base_path, num_designs)
             return
 
     _submit_batch_job(input_path, workers, queue, dry_run, base_path, after, num_designs)
@@ -343,7 +343,7 @@ def _submit_batch_job(
         raise SystemExit(1)
 
 
-def _run_local_mode(input_path: Path, base_path: str = BATCH_JOBS_BASE):
+def _run_local_mode(input_path: Path, base_path: str = BATCH_JOBS_BASE, num_designs: int = 1):
     """Run RFD3 locally in a Docker container after host-side fast-fail checks."""
     machine = platform.machine()
     if machine not in {"x86_64", "amd64"}:
@@ -395,6 +395,8 @@ def _run_local_mode(input_path: Path, base_path: str = BATCH_JOBS_BASE):
         f"FOUNDRY_CHECKPOINT_DIRS={CHECKPOINT_SUBTREE}",
         "-e",
         f"{OFFLINE_ENV_VAR}=1",
+        "-e",
+        f"NUM_DESIGNS={num_designs}",
         DEFAULT_IMAGE_URI,
     ]
 

{dh_cli-0.12.0 → dh_cli-0.12.2}/tests/test_rfd3_command.py

@@ -228,6 +228,22 @@ class TestLocalRun:
         assert "all" in cmd
         assert rfd3_cmd.DEFAULT_IMAGE_URI in cmd
 
+    def test_local_passes_num_designs(self, runner, input_dir):
+        """A local run with --num-designs N sets NUM_DESIGNS=N in the docker env."""
+        _write_config(input_dir, "design1")
+        result, run_mock = _run_local_directly(runner, input_dir, extra_args=["--num-designs", "4"])
+        assert result.exit_code == 0, result.output
+        cmd = run_mock.call_args[0][0]
+        assert "NUM_DESIGNS=4" in cmd
+
+    def test_local_default_num_designs_is_one(self, runner, input_dir):
+        """With no --num-designs flag, the local run defaults NUM_DESIGNS=1."""
+        _write_config(input_dir, "design1")
+        result, run_mock = _run_local_directly(runner, input_dir)
+        assert result.exit_code == 0, result.output
+        cmd = run_mock.call_args[0][0]
+        assert "NUM_DESIGNS=1" in cmd
+
 
 class TestLocalFastFail:
     def test_local_nonx86_fails_fast(self, runner, input_dir):