nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/cli/commands/object_store/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/cli/commands/object_store/object_store.py

@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import importlib
+import logging
+import mimetypes
+import time
+from pathlib import Path
+
+import click
+
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.object_store import ObjectStoreBaseConfig
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
+
+logger = logging.getLogger(__name__)
+
+STORE_CONFIGS = {
+    "s3": {
+        "module": "nat.plugins.s3.object_store", "config_class": "S3ObjectStoreClientConfig"
+    },
+    "mysql": {
+        "module": "nat.plugins.mysql.object_store", "config_class": "MySQLObjectStoreClientConfig"
+    },
+    "redis": {
+        "module": "nat.plugins.redis.object_store", "config_class": "RedisObjectStoreClientConfig"
+    }
+}
+
+
+def get_object_store_config(**kwargs) -> ObjectStoreBaseConfig:
+    """Process common object store arguments and return the config class"""
+    store_type = kwargs.pop("store_type")
+    config = STORE_CONFIGS[store_type]
+    module = importlib.import_module(config["module"])
+    config_class = getattr(module, config["config_class"])
+    return config_class(**kwargs)
+
+
+async def upload_file(object_store: ObjectStore, file_path: Path, key: str):
+    """
+    Upload a single file to object store.
+
+    Args:
+        object_store: The object store instance to use.
+        file_path: The path to the file to upload.
+        key: The key to upload the file to.
+    """
+    try:
+        data = await asyncio.to_thread(file_path.read_bytes)
+
+        item = ObjectStoreItem(data=data,
+                               content_type=mimetypes.guess_type(str(file_path))[0],
+                               metadata={
+                                   "original_filename": file_path.name,
+                                   "file_size": str(len(data)),
+                                   "file_extension": file_path.suffix,
+                                   "upload_timestamp": str(int(time.time()))
+                               })
+
+        # Upload using upsert to allow overwriting
+        await object_store.upsert_object(key, item)
+        click.echo(f"✅ Uploaded: {file_path.name} -> {key}")
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to upload {file_path.name}:\n{e}") from e
+
+
+def object_store_command_decorator(async_func):
+    """
+    Decorator that handles the common object store command pattern.
+
+    The decorated function should take (store: ObjectStore, kwargs) as parameters
+    and return an exit code (0 for success).
+    """
+
+    @click.pass_context
+    def wrapper(ctx: click.Context, **kwargs):
+        config = ctx.obj["store_config"]
+
+        async def work():
+            async with WorkflowBuilder() as builder:
+                await builder.add_object_store(name="store", config=config)
+                store = await builder.get_object_store_client("store")
+                return await async_func(store, **kwargs)
+
+        try:
+            exit_code = asyncio.run(work())
+        except Exception as e:
+            raise click.ClickException(f"Command failed: {e}") from e
+        if exit_code != 0:
+            raise click.ClickException(f"Command failed with exit code {exit_code}")
+        return exit_code
+
+    return wrapper
+
+
+@click.command(name="upload", help="Upload a directory to an object store.")
+@click.argument("local_dir",
+                type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
+                required=True)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def upload_command(store: ObjectStore, local_dir: Path, **_kwargs):
+    """
+    Upload a directory to an object store.
+
+    Args:
+        local_dir: The local directory to upload.
+        store: The object store to use.
+        _kwargs: Additional keyword arguments.
+    """
+    try:
+        click.echo(f"📁 Processing directory: {local_dir}")
+        file_count = 0
+
+        # Process each file recursively
+        for file_path in local_dir.rglob("*"):
+            if file_path.is_file():
+                key = file_path.relative_to(local_dir).as_posix()
+                await upload_file(store, file_path, key)
+                file_count += 1
+
+        click.echo(f"✅ Directory uploaded successfully! {file_count} files uploaded.")
+        return 0
+
+    except Exception as e:
+        raise click.ClickException(f"❌ Failed to upload directory {local_dir}:\n  {e}") from e
+
+
+@click.command(name="delete", help="Delete files from an object store.")
+@click.argument("keys", type=str, required=True, nargs=-1)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def delete_command(store: ObjectStore, keys: list[str], **_kwargs):
+    """
+    Delete files from an object store.
+
+    Args:
+        store: The object store to use.
+        keys: The keys to delete.
+        _kwargs: Additional keyword arguments.
+    """
+    deleted_count = 0
+    failed_count = 0
+    for key in keys:
+        try:
+            await store.delete_object(key)
+            click.echo(f"✅ Deleted: {key}")
+            deleted_count += 1
+        except Exception as e:
+            click.echo(f"❌ Failed to delete {key}: {e}")
+            failed_count += 1
+
+    click.echo(f"✅ Deletion completed! {deleted_count} keys deleted. {failed_count} keys failed to delete.")
+    return 0 if failed_count == 0 else 1
+
+
+@click.group(name="object-store", invoke_without_command=False, help="Manage object store operations.")
+def object_store_command(**_kwargs):
+    """Manage object store operations including uploading files and directories."""
+    pass
+
+
+def register_object_store_commands():
+
+    @click.group(name="s3", invoke_without_command=False, help="S3 object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--endpoint-url", type=str, help="S3 endpoint URL")
+    @click.option("--access-key", type=str, help="S3 access key")
+    @click.option("--secret-key", type=str, help="S3 secret key")
+    @click.option("--region", type=str, help="S3 region")
+    @click.pass_context
+    def s3(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="s3", **kwargs)
+
+    @click.group(name="mysql", invoke_without_command=False, help="MySQL object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="MySQL host")
+    @click.option("--port", type=int, help="MySQL port")
+    @click.option("--db", type=str, help="MySQL database name")
+    @click.option("--username", type=str, help="MySQL username")
+    @click.option("--password", type=str, help="MySQL password")
+    @click.pass_context
+    def mysql(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="mysql", **kwargs)
+
+    @click.group(name="redis", invoke_without_command=False, help="Redis object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="Redis host")
+    @click.option("--port", type=int, help="Redis port")
+    @click.option("--db", type=int, help="Redis db")
+    @click.pass_context
+    def redis(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="redis", **kwargs)
+
+    commands = {"s3": s3, "mysql": mysql, "redis": redis}
+
+    for store_type, config in STORE_CONFIGS.items():
+        try:
+            importlib.import_module(config["module"])
+            command = commands[store_type]
+            object_store_command.add_command(command, name=store_type)
+            command.add_command(upload_command, name="upload")
+            command.add_command(delete_command, name="delete")
+        except ImportError:
+            pass
+
+
+register_object_store_commands()

nat/cli/commands/optimize.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from pathlib import Path
+
+import click
+
+from nat.data_models.optimizer import OptimizerRunConfig
+from nat.profiler.parameter_optimization.optimizer_runtime import optimize_config
+
+logger = logging.getLogger(__name__)
+
+
+@click.group(name=__name__, invoke_without_command=True, help="Optimize a workflow with the specified dataset.")
+@click.option(
+    "--config_file",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    required=True,
+    help="A JSON/YAML file that sets the parameters for the workflow and evaluation.",
+)
+@click.option(
+    "--dataset",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    required=False,
+    help="A json file with questions and ground truth answers. This will override the dataset path in the config file.",
+)
+@click.option(
+    "--result_json_path",
+    type=str,
+    default="$",
+    help=("A JSON path to extract the result from the workflow. Use this when the workflow returns "
+          "multiple objects or a dictionary. For example, '$.output' will extract the 'output' field "
+          "from the result."),
+)
+@click.option(
+    "--endpoint",
+    type=str,
+    default=None,
+    help="Use endpoint for running the workflow. Example: http://localhost:8000/generate",
+)
+@click.option(
+    "--endpoint_timeout",
+    type=int,
+    default=300,
+    help="HTTP response timeout in seconds. Only relevant if endpoint is specified.",
+)
+@click.pass_context
+def optimizer_command(ctx, **kwargs) -> None:
+    """ Optimize workflow with the specified dataset"""
+    pass
+
+
+async def run_optimizer(config: OptimizerRunConfig):
+    await optimize_config(config)
+
+
+@optimizer_command.result_callback(replace=True)
+def run_optimizer_callback(
+    processors,  # pylint: disable=unused-argument
+    *,
+    config_file: Path,
+    dataset: Path,
+    result_json_path: str,
+    endpoint: str,
+    endpoint_timeout: int,
+):
+    """Run the optimizer with the provided config file and dataset."""
+    config = OptimizerRunConfig(
+        config_file=config_file,
+        dataset=dataset,
+        result_json_path=result_json_path,
+        endpoint=endpoint,
+        endpoint_timeout=endpoint_timeout,
+    )
+
+    asyncio.run(run_optimizer(config))

nat/cli/commands/registry/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/cli/commands/registry/publish.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from contextlib import AsyncExitStack
+from pathlib import Path
+
+import click
+
+from nat.data_models.registry_handler import RegistryHandlerBaseConfig
+from nat.utils.data_models.schema_validator import validate_yaml
+
+logger = logging.getLogger(__name__)
+
+
+async def publish_artifact(registry_handler_config: RegistryHandlerBaseConfig, package_root: str) -> None:
+
+    from nat.cli.type_registry import GlobalTypeRegistry
+    from nat.registry_handlers.package_utils import build_artifact
+
+    registry = GlobalTypeRegistry.get()
+
+    async with AsyncExitStack() as stack:
+
+        registry_handler_info = registry.get_registry_handler(type(registry_handler_config))
+        registry_handler = await stack.enter_async_context(registry_handler_info.build_fn(registry_handler_config))
+        try:
+            artifact = build_artifact(package_root=package_root)
+        except Exception as e:
+            logger.exception("Error building artifact: %s", e)
+            return
+        await stack.enter_async_context(registry_handler.publish(artifact=artifact))
+
+
+@click.group(name=__name__,
+             invoke_without_command=True,
+             help=("Publish local NAT artifacts to a remote "
+                   "registry from package repository."))
+@click.option(
+    "--config_file",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    callback=validate_yaml,
+    required=False,
+    help=("A YAML file to override configured channel settings."),
+)
+@click.option(
+    "-c",
+    "--channel",
+    type=str,
+    required=True,
+    help=("The remote registry channel to use when publishing the NAT artifact."),
+)
+@click.argument("package_root", type=str)
+def publish(channel: str, config_file: str, package_root: str) -> None:
+    """
+    Publish NAT artifacts with the specified configuration
+    """
+    from nat.settings.global_settings import GlobalSettings
+
+    settings = GlobalSettings().get()
+
+    if (config_file is not None):
+        settings = settings.override_settings(config_file)
+
+    try:
+        publish_channel_config = settings.channels.get(channel)
+
+        if (publish_channel_config is None):
+            logger.error("Publish channel '%s' has not been configured.", channel)
+            return
+    except Exception as e:
+        logger.exception("Error loading user settings: %s", e)
+        return
+
+    asyncio.run(publish_artifact(registry_handler_config=publish_channel_config, package_root=package_root))

nat/cli/commands/registry/pull.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from contextlib import AsyncExitStack
+from pathlib import Path
+
+import click
+
+from nat.data_models.registry_handler import RegistryHandlerBaseConfig
+from nat.utils.data_models.schema_validator import validate_yaml
+
+logger = logging.getLogger(__name__)
+
+
+async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, packages: list[str]) -> None:
+
+    from nat.cli.type_registry import GlobalTypeRegistry
+    from nat.registry_handlers.schemas.package import PackageNameVersion
+    from nat.registry_handlers.schemas.pull import PullPackageWhl
+    from nat.registry_handlers.schemas.pull import PullRequestPackages
+
+    registry = GlobalTypeRegistry.get()
+
+    async with AsyncExitStack() as stack:
+
+        registry_handler_info = registry.get_registry_handler(type(registry_handler_config))
+        registry_handler = await stack.enter_async_context(registry_handler_info.build_fn(registry_handler_config))
+
+        try:
+            package_list = []
+            for package in packages:
+
+                package_data = {}
+
+                assert len(package) > 0, f"Supplied invalid package '{package}'."
+
+                if package[:-4] == ".whl":
+                    package_data["whl_path"] = package
+                    package_list.append(PullPackageWhl(**package_data))
+                else:
+                    package_split = package.split("==")
+
+                    assert len(package_split) in (1, 2), f"Supplied invalid package '{package}'."
+
+                    package_data["name"] = package_split[0]
+
+                    if (package_split == 2):
+                        package_data["version"] = package_split[1]
+
+                    package_list.append(PackageNameVersion(**package_data))
+
+            validated_packages = PullRequestPackages(packages=package_list)
+
+        except Exception as e:
+            logger.exception("Error processing package names: %s", e)
+            return
+
+        await stack.enter_async_context(registry_handler.pull(packages=validated_packages))
+
+
+@click.group(name=__name__,
+             invoke_without_command=True,
+             help=("Pull NAT artifacts from a remote registry "
+                   "by package name."))
+@click.option(
+    "--config_file",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    callback=validate_yaml,
+    required=False,
+    help=("A YAML file to override the channel settings."),
+)
+@click.option(
+    "-c",
+    "--channel",
+    type=str,
+    required=True,
+    help=("The remote registry channel to use when pulling the NAT artifact."),
+)
+@click.argument("packages", type=str)
+def pull(channel: str, config_file: str, packages: str) -> None:
+    """
+    Pull NAT artifacts from a remote registry channel.
+    """
+
+    from nat.settings.global_settings import GlobalSettings
+
+    packages = packages.split()
+
+    settings = GlobalSettings().get()
+
+    if (config_file is not None):
+        settings = settings.override_settings(config_file)
+
+    try:
+        pull_channel_config = settings.channels.get(channel)
+
+        if (pull_channel_config is None):
+            logger.error("Pull channel '%s' has not been configured.", channel)
+            return
+    except Exception as e:
+        logger.exception("Error loading user settings: %s", e)
+        return
+
+    asyncio.run(pull_artifact(pull_channel_config, packages))

nat/cli/commands/registry/registry.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import click
+
+from nat.cli.commands.registry.publish import publish
+from nat.cli.commands.registry.pull import pull
+from nat.cli.commands.registry.remove import remove
+from nat.cli.commands.registry.search import search
+
+logger = logging.getLogger(__name__)
+
+
+@click.group(name=__name__, invoke_without_command=False, help="Utility to configure NAT remote registry channels.")
+def registry_command(**kwargs):
+    pass
+
+
+registry_command.add_command(publish, "publish")
+registry_command.add_command(pull, "pull")
+registry_command.add_command(remove, "remove")
+registry_command.add_command(search, "search")