feat!: Support adaptive chunking, general CLI improvement

README.md

@@ -22,13 +22,13 @@ This library requires:
 
 ## Versioning
 
-This package's version follows [Semantic Versioning 2.0](https://semver.org/), but is still considered to be in its 
+This package's version follows [Semantic Versioning 2.0](https://semver.org/), but is still considered to be in its
 initial development, thus backwards incompatible versions are denoted by minor version bumps. To help illustrate how
 versions will increment during this initial development stage, they are described below:
 
-1. The MAJOR version is currently 0, indicating initial development. 
-2. The MINOR version is currently incremented when backwards incompatible changes are introduced to the public API. 
-3. The PATCH version is currently incremented when bug fixes or backwards compatible changes are introduced to the public API. 
+1. The MAJOR version is currently 0, indicating initial development.
+2. The MINOR version is currently incremented when backwards incompatible changes are introduced to the public API.
+3. The PATCH version is currently incremented when bug fixes or backwards compatible changes are introduced to the public API.
 
 ## Contributing
 
@@ -64,7 +64,7 @@ Template at '/path/to/job.template.json' passes validation checks!
 
 ### `summary`
 
-Displays summary information about a sample Job or Step, and the Steps and Tasks therein. The user may provide parameters to 
+Displays summary information about a sample Job or Step, and the Steps and Tasks therein. The user may provide parameters to
 customize the Job, as parameters can have an impact on the amount of Steps and Tasks that a job consists of.
 
 #### Arguments
@@ -145,7 +145,7 @@ Session ended successfully
 Job: MyJob
 Step: Step1
 Duration: 1.0 seconds
-Tasks run: 1
+Chunks run: 1
 
 ```
 
@@ -181,11 +181,11 @@ See [VERIFYING_PGP_SIGNATURE](VERIFYING_PGP_SIGNATURE.md) for more information.
 
 ## Security
 
-We take all security reports seriously. When we receive such reports, we will 
-investigate and subsequently address any potential vulnerabilities as quickly 
-as possible. If you discover a potential security issue in this project, please 
+We take all security reports seriously. When we receive such reports, we will
+investigate and subsequently address any potential vulnerabilities as quickly
+as possible. If you discover a potential security issue in this project, please
 notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/)
-or directly via email to [AWS Security]([email protected]). Please do not 
+or directly via email to [AWS Security]([email protected]). Please do not
 create a public GitHub issue in this project.
 
 ## License

pyproject.toml

@@ -29,8 +29,8 @@ classifiers = [
   "Intended Audience :: End Users/Desktop"
 ]
 dependencies = [
-  "openjd-sessions == 0.10.*",
-  "openjd-model == 0.6.*"
+  "openjd-sessions >= 0.10.1,< 0.11",
+  "openjd-model == 0.7.*"
 ]
 
 [project.urls]

src/openjd/__main__.py

@@ -1,26 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-import sys
-import traceback
+from .cli._create_argparser import main
 
-from .cli._create_argparser import create_argparser
 
-__all__ = ("main",)
-
-
-def main() -> None:
-    parser = create_argparser()
-
-    args = parser.parse_args(sys.argv[1:])
-    try:
-        # Raises:
-        #  SystemExit - on failure
-        args.func(args)
-    except Exception as exc:
-        print(f"ERROR: {str(exc)}", file=sys.stderr)
-        traceback.print_exc()
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()
+main()

src/openjd/cli/__init__.py

@@ -1 +1,5 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+from ._create_argparser import main
+
+__all__ = ["main"]

src/openjd/cli/_check/__init__.py

@@ -1,7 +1,7 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
 from .._common import add_common_arguments, CommonArgument, SubparserGroup
-from ._check_command import do_check
+from ._check_command import add_check_arguments, do_check
 
 
 def populate_argparser(subcommands: SubparserGroup) -> None:
@@ -12,7 +12,7 @@ def populate_argparser(subcommands: SubparserGroup) -> None:
         description="Given an Open Job Description template file, parse the file and run validation checks against it to ensure that it is correctly formed.",
     )
 
-    # `check` has no unique arguments;
     # add all arguments through `add_common_arguments`
     add_common_arguments(check_parser, {CommonArgument.PATH})
+    add_check_arguments(check_parser)
     check_parser.set_defaults(func=do_check)

src/openjd/cli/_check/_check_command.py

@@ -1,20 +1,32 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-from argparse import Namespace
+from argparse import ArgumentParser, Namespace
 from openjd.model import (
     DecodeValidationError,
     TemplateSpecificationVersion,
     decode_job_template,
     decode_environment_template,
 )
 
-from .._common import read_template, OpenJDCliResult, print_cli_result
+from .._common import (
+    add_extensions_argument,
+    read_template,
+    OpenJDCliResult,
+    print_cli_result,
+    process_extensions_argument,
+)
+
+
+def add_check_arguments(run_parser: ArgumentParser):
+    add_extensions_argument(run_parser)
 
 
 @print_cli_result
 def do_check(args: Namespace) -> OpenJDCliResult:
     """Open a provided template file and check its schema for errors."""
 
+    extensions = process_extensions_argument(args.extensions)
+
     try:
         # Raises: RuntimeError
         template_object = read_template(args.path)
@@ -27,7 +39,7 @@ def do_check(args: Namespace) -> OpenJDCliResult:
 
         # Raises: DecodeValidationError
         if TemplateSpecificationVersion.is_job_template(template_version):
-            decode_job_template(template=template_object, supported_extensions=["TASK_CHUNKING"])
+            decode_job_template(template=template_object, supported_extensions=extensions)
         elif TemplateSpecificationVersion.is_environment_template(template_version):
             decode_environment_template(template=template_object)
         else:

src/openjd/cli/_common/__init__.py

@@ -9,6 +9,7 @@
 import yaml
 import os
 
+from ._extensions import add_extensions_argument, process_extensions_argument
 from ._job_from_template import (
     job_from_template,
     get_job_params,
@@ -23,9 +24,11 @@
 from openjd.model import DecodeValidationError, Job
 
 __all__ = [
+    "add_extensions_argument",
     "get_doc_type",
     "get_job_params",
     "get_params_from_file",
+    "process_extensions_argument",
     "read_template",
     "read_job_template",
     "read_environment_template",
@@ -109,10 +112,10 @@ def add(self, name: str, description: str, **kwargs) -> ArgumentParser:
         return self.group.add_parser(name, **kwargs)
 
 
-def generate_job(args: Namespace) -> Job:
+def generate_job(args: Namespace, *, supported_extensions: list[str]) -> Job:
     try:
         # Raises: RuntimeError, DecodeValidationError
-        template = read_job_template(args.path)
+        template = read_job_template(args.path, supported_extensions=supported_extensions)
         # Raises: RuntimeError
         return job_from_template(
             template,

src/openjd/cli/_common/_extensions.py

@@ -0,0 +1,34 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+from argparse import ArgumentParser
+from typing import Optional
+
+# This is the list of Open Job Description extensions with implemented support
+SUPPORTED_EXTENSIONS = ["TASK_CHUNKING"]
+
+
+def add_extensions_argument(run_parser: ArgumentParser):
+    run_parser.add_argument(
+        "--extensions",
+        help=f"A comma-separated list of Open Job Description extension names to enable. Defaults to all that are implemented: {','.join(SUPPORTED_EXTENSIONS)}",
+    )
+
+
+def process_extensions_argument(extensions: Optional[str]) -> list[str]:
+    """Process the comma-separated extensions argument and return a list of supported extensions."""
+
+    # If the option is not provided, default to all the supported extensions.
+    if extensions is None:
+        return SUPPORTED_EXTENSIONS
+
+    extensions_list = [
+        extension.strip().upper() for extension in extensions.split(",") if extension.strip() != ""
+    ]
+
+    unsupported_extensions = set(extensions_list) - set(SUPPORTED_EXTENSIONS)
+    if unsupported_extensions:
+        raise ValueError(
+            f"Unsupported Open Job Description extension(s): {', '.join(sorted(unsupported_extensions))}"
+        )
+
+    return extensions_list

src/openjd/cli/_common/_validation_utils.py

@@ -53,7 +53,7 @@ def read_template(template_file: Path) -> dict[str, Any]:
     return template_object
 
 
-def read_job_template(template_file: Path) -> JobTemplate:
+def read_job_template(template_file: Path, *, supported_extensions: list[str]) -> JobTemplate:
     """Open a JSON or YAML-formatted file and attempt to parse it into a JobTemplate object.
     Raises a RuntimeError if the file doesn't exist or can't be opened, and raises a
     DecodeValidationError if its contents can't be parsed into a valid JobTemplate.
@@ -62,7 +62,9 @@ def read_job_template(template_file: Path) -> JobTemplate:
     template_object = read_template(template_file)
 
     # Raises: DecodeValidationError
-    template = decode_job_template(template=template_object, supported_extensions=["TASK_CHUNKING"])
+    template = decode_job_template(
+        template=template_object, supported_extensions=supported_extensions
+    )
 
     return template
 

src/openjd/cli/_create_argparser.py

@@ -1,6 +1,9 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
 from argparse import ArgumentParser
+import sys
+import traceback
+from typing import Optional
 
 from ._common import SubparserGroup
 
@@ -28,3 +31,21 @@ def create_argparser() -> ArgumentParser:
     populate_run_subparser(subcommands)
     populate_schema_subparser(subcommands)
     return parser
+
+
+def main(arg_list: Optional[list[str]] = None) -> None:
+    """Main function for invoking the CLI"""
+    parser = create_argparser()
+
+    if arg_list is None:
+        arg_list = sys.argv[1:]
+
+    args = parser.parse_args(arg_list)
+    try:
+        # Raises:
+        #  SystemExit - on failure
+        args.func(args)
+    except Exception as exc:
+        print(f"ERROR: {str(exc)}", file=sys.stderr)
+        traceback.print_exc()
+        sys.exit(1)

src/openjd/cli/_run/__init__.py

@@ -8,7 +8,7 @@ def populate_argparser(subcommands: SubparserGroup) -> None:
     """Adds the `run` command and all of its arguments to the given parser."""
     run_parser = subcommands.add(
         "run",
-        description="Takes a Job Template and Step name, then runs Tasks from that Step.",
+        description="Takes a Job Template and runs the entire job or a selected Step from the job.",
         usage="openjd run JOB_TEMPLATE_PATH [arguments]",
     )
     add_common_arguments(run_parser, {CommonArgument.PATH, CommonArgument.JOB_PARAMS})

src/openjd/cli/_run/_local_session/_actions.py

@@ -1,10 +1,27 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
+from enum import Enum
+
 from openjd.model import Step, TaskParameterSet
 from openjd.model.v2023_09 import Environment
 from openjd.sessions import Session
 
 
+class EnvironmentType(str, Enum):
+    """
+    The three different types of environment types that can be entered/exited in a session.
+    """
+
+    EXTERNAL = "EXTERNAL"
+    JOB = "JOB"
+    STEP = "STEP"
+    ALL = "ALL"
+
+    def matches(self, other: "EnvironmentType") -> bool:
+        """Environment types match if they are equal, or one of them is ALL."""
+        return self == other or self == EnvironmentType.ALL or other == EnvironmentType.ALL
+
+
 class SessionAction:
     _session: Session
     duration: float
@@ -56,13 +73,17 @@ def __str__(self):
 
 class ExitEnvironmentAction(SessionAction):
     _id: str
+    _keep_session_running: bool
 
-    def __init__(self, session: Session, id: str):
+    def __init__(self, session: Session, id: str, keep_session_running: bool):
         super(ExitEnvironmentAction, self).__init__(session)
         self._id = id
+        self._keep_session_running = keep_session_running
 
     def run(self):
-        self._session.exit_environment(identifier=self._id)
+        self._session.exit_environment(
+            identifier=self._id, keep_session_running=self._keep_session_running
+        )
 
     def __str__(self):
         return f"Exit Environment '{self._id}'"