Spaces:
Running on L40S
Running on L40S
File size: 3,505 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
from pathlib import Path
from typing import Annotated
import pydantic
import requests
import tyro
from cosmos_framework.inference.args import OmniSampleOverrides
from cosmos_framework.inference.common.args import ResolvedPath, SampleOutputs, tyro_cli
from cosmos_framework.utils import log
class Args(pydantic.BaseModel):
output_dir: Annotated[ResolvedPath, tyro.conf.arg(aliases=("-o",))]
"""Output directory."""
input_files: Annotated[list[ResolvedPath], tyro.conf.arg(aliases=("-i",))]
"""Path to the inference parameter file(s).
If multiple files are provided, the model will be loaded once and all the samples will be run sequentially.
Accepts glob patterns (e.g. `inputs/*.json`).
"""
overrides: OmniSampleOverrides = pydantic.Field(default_factory=OmniSampleOverrides)
"""Overrides for the inference parameters."""
host: str = "localhost"
"""The hostname of the Ray Serve endpoint."""
port: int = 8000
"""The port of the Ray Serve endpoint."""
timeout: int = 600
"""The timeout for the request in seconds."""
@property
def server(self) -> str:
return f"http://{self.host}:{self.port}"
def handle_response(response: requests.Response) -> None:
try:
response.raise_for_status()
except requests.exceptions.HTTPError as e:
try:
detail = response.json().get("detail")
except Exception:
detail = None
if detail:
raise ValueError(f"{detail}") from e
raise e
def submit(args: Args):
args.output_dir.mkdir(parents=True, exist_ok=True)
sample_args_list = OmniSampleOverrides.from_files(args.input_files, overrides=args.overrides)
for i_sample, sample_args in enumerate(sample_args_list):
assert sample_args.name
log.info(f"[{i_sample + 1}/{len(sample_args_list)}] Submitting sample '{sample_args.name}'")
log.debug(f"{sample_args.__class__.__name__}({sample_args})")
sample_dir = args.output_dir / sample_args.name
sample_dir.mkdir(parents=True, exist_ok=True)
sample_args_file = sample_dir / "sample_args.json"
sample_args_file.write_text(sample_args.model_dump_json())
log.info(f"Saved sample args to '{sample_args_file}'")
response = requests.post(
f"{args.server}/generate", json=sample_args.model_dump(mode="json"), timeout=args.timeout
)
handle_response(response)
response_dict = response.json()
sample_outputs = SampleOutputs.model_validate(response_dict)
# Download files
def download_file(remote_file: Path) -> Path:
response = requests.get(f"{args.server}/outputs/{remote_file}", timeout=args.timeout)
handle_response(response)
local_file = sample_dir.joinpath(*remote_file.parts[1:])
local_file.write_bytes(response.content)
return local_file
sample_outputs = sample_outputs.map_files(download_file)
sample_outputs_file = sample_dir / "sample_outputs.json"
sample_outputs_file.write_text(sample_outputs.model_dump_json())
log.info(f"Saved sample outputs to '{sample_outputs_file}'")
def main():
args = tyro_cli(Args, description=__doc__, config=(tyro.conf.OmitArgPrefixes,))
submit(args)
if __name__ == "__main__":
main()
|