You are viewing an outdated version of the documentation.

This documentation is for an older version (1.4.7) of Dagster. You can view the version of this page from our latest release below.

Source code for dagster_wandb.launch.ops

from dagster import OpExecutionContext, op
from wandb.sdk.launch import launch
from wandb.sdk.launch.launch_add import launch_add

from .configs import launch_agent_config, launch_config


def raise_on_invalid_config(context: OpExecutionContext):
    entity = context.resources.wandb_config["entity"]
    if entity == "":
        raise RuntimeError(
            "(dagster_wandb) An empty string was provided for the 'entity' property of the"
            " 'wandb_config'."
        )

    project = context.resources.wandb_config["project"]
    if project == "":
        raise RuntimeError(
            "(dagster_wandb) An empty string was provided for the 'project' property of the"
            " 'wandb_config'."
        )


[docs]@op( required_resource_keys={"wandb_resource", "wandb_config"}, config_schema=launch_agent_config(), ) def run_launch_agent(context: OpExecutionContext): """It starts a Launch Agent and runs it as a long running process until stopped manually. Agents are processes that poll launch queues and execute the jobs (or dispatch them to external services to be executed) in order. **Example:** .. code-block:: YAML # config.yaml resources: wandb_config: config: entity: my_entity project: my_project ops: run_launch_agent: config: max_jobs: -1 queues: - my_dagster_queue .. code-block:: python from dagster_wandb.launch.ops import run_launch_agent from dagster_wandb.resources import wandb_resource from dagster import job, make_values_resource @job( resource_defs={ "wandb_config": make_values_resource( entity=str, project=str, ), "wandb_resource": wandb_resource.configured( {"api_key": {"env": "WANDB_API_KEY"}} ), }, ) def run_launch_agent_example(): run_launch_agent() """ raise_on_invalid_config(context) config = { "entity": context.resources.wandb_config["entity"], "project": context.resources.wandb_config["project"], **context.op_config, } context.log.info(f"Launch agent configuration: {config}") context.log.info("Running Launch agent...") launch.create_and_run_agent(api=context.resources.wandb_resource["api"], config=config)
[docs]@op( required_resource_keys={ "wandb_resource", "wandb_config", }, config_schema=launch_config(), ) def run_launch_job(context: OpExecutionContext): """Executes a Launch job. A Launch job is assigned to a queue in order to be executed. You can create a queue or use the default one. Make sure you have an active agent listening to that queue. You can run an agent inside your Dagster instance but can also consider using a deployable agent in Kubernetes. **Example:** .. code-block:: YAML # config.yaml resources: wandb_config: config: entity: my_entity project: my_project ops: my_launched_job: config: entry_point: - python - train.py queue: my_dagster_queue uri: https://github.com/wandb/example-dagster-integration-with-launch .. code-block:: python from dagster_wandb.launch.ops import run_launch_job from dagster_wandb.resources import wandb_resource from dagster import job, make_values_resource @job( resource_defs={ "wandb_config": make_values_resource( entity=str, project=str, ), "wandb_resource": wandb_resource.configured( {"api_key": {"env": "WANDB_API_KEY"}} ), }, ) def run_launch_job_example(): run_launch_job.alias("my_launched_job")() # we rename the job with an alias """ raise_on_invalid_config(context) config = { "entity": context.resources.wandb_config["entity"], "project": context.resources.wandb_config["project"], **context.op_config, } context.log.info(f"Launch job configuration: {config}") queue = context.op_config.get("queue") if queue is None: context.log.info("No queue provided, running Launch job locally") launch.run(api=context.resources.wandb_resource["api"], config=config) else: synchronous = config.get("synchronous", True) config.pop("synchronous", None) queued_run = launch_add(**config) if synchronous is True: context.log.info( f"Synchronous Launch job added to queue with name={queue}. Waiting for" " completion..." ) queued_run.wait_until_finished() else: context.log.info(f"Asynchronous Launch job added to queue with name={queue}")