Skip to content

rats.aml

Submit applications as aml jobs, sending necessary context from the driver node.

__all__ = ['AmlEnvironment', 'AmlIO', 'AmlJobContext', 'AmlJobDetails', 'AmlWorkspace', 'AppConfigs', 'AppServices', 'Application', 'main', 'submit'] module-attribute

AppConfigs

Configuration services and groups used by the rats.aml.Application.

Some of these config services and groups alter the construction of the services defined in AppServices.

JOB_DETAILS = apps.ServiceId[AmlJobDetails]('job-details.config') class-attribute instance-attribute

Aml job runtime information containing compute, input/output, workspace, and environment info.

This service defers to more specific configurations that can be defined individually if you don't require providing the entire configuration.

class Plugin(apps.Container):

    @apps.service(AppConfigs.JOB_DETAILS)
    def _runtime(self) -> AmlConfig:
        AmlConfig(
            compute="[compute-cluster-name]",
            inputs={},
            outputs={},
            workspace=self._app.get(AppConfigs.WORKSPACE),
            environment=self._app.get(AppConfigs.ENVIRONMENT),
        )

COMPUTE = apps.ServiceId[str]('compute.config') class-attribute instance-attribute

The name of the compute cluster the aml job should be submitted to.

This can be a full resource id in cases where the cluster is part of a workspace other than the one being submitted to.

ENVIRONMENT = apps.ServiceId[AmlEnvironment]('environment.config') class-attribute instance-attribute

The workspace environment that should be created and used for the aml job.

If the environment already exists, it will be updated before submitting the job. We recommend using an environment that leverages a pre-built container image and not relying on aml to upload your code and build a new container on each job submission. This allows you to rely on images that you've ideally built and tested in CI pipelines without modification.

WORKSPACE = apps.ServiceId[AmlWorkspace]('workspace.config') class-attribute instance-attribute

The workspace information for submitting aml jobs into.

COMMAND_KWARGS = apps.ServiceId[Mapping[str, Any]]('command-kwargs.config') class-attribute instance-attribute

Additional arguments to be added to the azure.ai.ml.command call when submitting a job.

In cases where the rats.aml module does not provide the desired configuration for a job, this service can be used to add a dictionary of kwargs that will be added to the command construction function without validation.

class Plugin(apps.Container):
    @apps.service(AppConfigs.COMMAND_KWARGS)
    def _ml_command_kwargs(self) -> Mapping[str, Any]:
        return {
            "resources": {"instance_count": 1},
        }

JOB_CONTEXT = apps.ServiceId[AmlJobContext]('job-context.config') class-attribute instance-attribute

Context element containing basic job information and always shared with the worker node.

This context object is always added to any registered rats.aml.AppConfigs.CONTEXT_COLLECTION in order to provide a handful of details that might be useful for tracking larger pipelines.

INPUTS = apps.ServiceId[Mapping[str, AmlIO]]('inputs.config-group') class-attribute instance-attribute

Aml inputs used by the job and mounted on the worker nodes.

On the worker node, the path to the mounted input directory is exposed as an environment variable prefixed with RATS_AML_PATH_, followed by the uppercase [input-name]. So an input registered like below, and named model_input will have an environment variable RATS_AML_PATH_MODEL_INPUT containing the local path to the mounted directory/file.

class Plugin(apps.Container):
    @apps.fallback_group(AppConfigs.INPUTS)
    def _inputs(self) -> Iterator[Mapping[str, AmlIO]]:
        from azure.ai.ml.constants import AssetTypes, InputOutputModes

        yield {
            f"[input-name]": AmlIO(
                type=AssetTypes.URI_FOLDER,
                path=f"abfss://[container]@[storage-account].dfs.core.windows.net/[path]",
                mode=InputOutputModes.RW_MOUNT,
            ),
        }

OUTPUTS = apps.ServiceId[Mapping[str, AmlIO]]('outputs.config-group') class-attribute instance-attribute

Aml outputs used by the job and mounted on the worker nodes.

On the worker node, the path to the mounted output directory is exposed as an environment variable prefixed with RATS_AML_PATH_, followed by the uppercase [output-name]. So an output registered like below, and named model_output will have an environment variable RATS_AML_PATH_MODEL_OUTPUT containing the local path to the mounted directory/file.

class Plugin(apps.Container):
    @apps.fallback_group(AppConfigs.OUTPUTS)
    def _outputs(self) -> Iterator[Mapping[str, AmlIO]]:
        from azure.ai.ml.constants import AssetTypes, InputOutputModes

        yield {
            f"[output-name]": AmlIO(
                type=AssetTypes.URI_FOLDER,
                path=f"abfss://[container]@[storage-account].dfs.core.windows.net/[path]",
                mode=InputOutputModes.RW_MOUNT,
            ),
        }

CONTEXT_COLLECTION = apps.ServiceId[app_context.Collection[Any]]('app-ctx-collection.config') class-attribute instance-attribute

Collection containing the entire context being sent to the aml job.

You should not set this service config directly and instead register to the rats.aml.AppConfigs.APP_CONTEXT service group to add elements to this collection. You can retrieve this service config on the aml job instance to retrieve the registered contexts.

class Application(apps.AppContainer):
    def execute(self) -> None:
        context_collection = self._app.get(aml.AppConfigs.CONTEXT_COLLECTION)
        print("loaded context:")
        for item in context_collection.items:
            print(f"{item.service_id} -> {item.values}")

CLI_ENVS = apps.ServiceId[Mapping[str, str]]('cli-envs.config-group') class-attribute instance-attribute

Dictionaries to merge and attach as environment variables to the submitted aml job.

The special environment variables like RATS_AML_RUN_CONTEXT are not affected by the definition of these service groups.

class Plugin(apps.Container):

    @apps.service_group(AppConfigs.CLI_ENVS)
    def _envs(self) -> Iterator[Mapping[str, str]]:
        yield {"MY_JOB_ENV": "some-useful-information"}

CLI_CWD = apps.ServiceId[str]('cli-cwd.config') class-attribute instance-attribute

The directory within the aml job container image from where to run the rats-aml run command.

This defaults to the directory of the component rats-aml is being run from, assuming the container has been built with this component at /opt/<project>/<component>, as is done in the Containerfile's in this repository. Define this config service if this convention doesn't match your project's.

APP_CONTEXT = apps.ServiceId[app_context.Context[Any]]('app-context.config-group') class-attribute instance-attribute

Service group containing context services to attach to the submitted aml job.

By default, rats.aml.AppConfigs.JOB_CONTEXT is attached and provides some basic information about the submitted job along with a unique id that can be used for tracking runs.

Application(app)

Bases: apps.AppContainer, cli.Container, apps.PluginMixin

CLI application for submitting rats applications to be run on aml.

Source code in rats/apps/_app_containers.py 
94
95
def __init__(self, app: Container) -> None:
    self._app = app

execute()

Runs the rats-aml cli that provides methods for listing and submitting aml jobs.

Source code in src/rats/aml/_app.py 
264
265
266
267
268
269
270
271
272
273
def execute(self) -> None:
    """Runs the `rats-aml` cli that provides methods for listing and submitting aml jobs."""
    argv = self._app.get(_CliContext.SERVICE_ID)
    cli.create_group(click.Group("rats-aml"), self).main(
        args=argv[1:],
        prog_name=Path(argv[0]).name,
        auto_envvar_prefix="RATS_AML",
        # don't end the process
        standalone_mode=False,
    )

AppServices

Service classes used and exposed by rats.aml.Application.

Many of the constructor arguments for these services are configurable by providing config services and groups defined in AppConfigs. Any service that isn't configurable to the extent needed can typically be overwritten by the user in a plugin container.

AML_CLIENT = apps.ServiceId['MLClient']('aml-client') class-attribute instance-attribute

Instance of the azure.ai.ml.MLClient used to submit this application's aml jobs.

AML_ENVIRONMENT_OPS = apps.ServiceId['EnvironmentOperations']('aml-environment-ops') class-attribute instance-attribute

azure.ai.ml.operations.EnvironmentOperations used to create and update the aml environment.

This service is not overwritable because it comes directly from the aml client defined by AppServices.AML_CLIENT.

AML_JOB_OPS = apps.ServiceId['JobOperations']('aml-job-ops') class-attribute instance-attribute

azure.ai.ml.operations.JobOperations used to submit and monitor the aml job for completion.

This service is not overwritable because it comes directly from the aml client defined by AppServices.AML_CLIENT.

AML_IDENTITY = apps.ServiceId['TokenCredential']('identity') class-attribute instance-attribute

azure.core.credentials.TokenCredential used by the MLClient.

By default, we use a DefaultAzureCredential, but a more specific one can be used in more advanced projects.

AmlEnvironment

Bases: NamedTuple

The details of the AML Environment being used by an aml job.

name instance-attribute

The name of the built environment in AML.

Often, this is the name of the component since rats typically has a 1:1 relationship between components and docker images. These environments are discoverable in the AML workspace and each environment can have many versions.

image instance-attribute

The full name, with registry and tag, of the container image for this environment.

In the AML workspace, this is shown as the "Docker image" and is sometimes used as the parent image when asking AML to build the environment for us. However, in rats, we typically build a full image using the standard tools like Docker, and have AML use the built images without making any runtime modifications to it. This ensures the images we build in CI pipelines are used without modifications after they have been tested (by you).

version instance-attribute

The version of the environment being used for AML jobs.

Typically, this matches the tag portion of the built container images. You will find a version drop-down when viewing the environment in your AML workspace.

full_name property

The AML environment and version.

Unlike the AmlEnvironment.image property, this is the full name of the AML environment, and not the container image. In most cases, this is the name of the component, and the container tag, but does not contain the registry information. This is a unique identifier within your AML workspace, but not globally, like container images.

AmlIO

Bases: NamedTuple

type instance-attribute

path instance-attribute

mode instance-attribute

AmlJobContext(uuid, job_details) dataclass

Context added by default to all aml jobs with values used by rats.aml.

uuid instance-attribute

A unique id tied to a given submission of an aml job.

job_details instance-attribute

The compute, workspace, environment, and input/outputs for the aml job.

AmlJobDetails

Bases: NamedTuple

compute instance-attribute

inputs instance-attribute

outputs instance-attribute

workspace instance-attribute

environment instance-attribute

AmlWorkspace

Bases: NamedTuple

The aml workspace used for submitted aml jobs.

subscription_id instance-attribute

Azure subscription containing the desired workspace.

resource_group_name instance-attribute

Azure resource group containing the desired workspace.

workspace_name instance-attribute

Azure aml workspace name jobs should be submitted to.

main()

Main entry-point to the rats-aml cli command.

Source code in src/rats/aml/_app.py 
574
575
576
577
def main() -> None:
    """Main entry-point to the `rats-aml` cli command."""
    apps.run_plugin(logs.ConfigureApplication)
    apps.run_plugin(Application)

submit(*app_ids, context=app_context.EMPTY_COLLECTION, wait=False)

Submit an AML job programmatically instead of calling rats-aml submit.

Parameters:

Name Type Description Default
app_ids str

list of the application to run on the remote aml job as found in pyproject.toml

 ()
context app_context.Collection[Any]

context to send to the remote aml job

 app_context.EMPTY_COLLECTION
wait bool

wait for the successful completion of the submitted aml job.

 False
Source code in src/rats/aml/_app.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
def submit(
    *app_ids: str,
    context: app_context.Collection[Any] = app_context.EMPTY_COLLECTION,
    wait: bool = False,
) -> None:
    """
    Submit an AML job programmatically instead of calling `rats-aml submit`.

    Args:
        app_ids: list of the application to run on the remote aml job as found in pyproject.toml
        context: context to send to the remote aml job
        wait: wait for the successful completion of the submitted aml job.
    """
    w = ["--wait"] if wait else []
    cmd = (
        "rats-aml",
        "submit",
        *app_ids,
        "--context",
        app_context.dumps(context),
        *w,
    )
    submitter = apps.AppBundle(
        app_plugin=Application,
        context=_CliContext(cmd),
    )
    submitter.execute()